llama : prune (#30565)

* llama : remove deprecated `llama_set_warmup`

Assisted-by: dsh:llama.cpp/ggml-org/DeepSeek-V4.1-Flash-1017-GGUF (Mac Studio, M5 Ultra, 256GB)

* llama : remove remaining deprecated API

Drop the APIs marked DEPRECATED in llama.h, the deprecated
llama_context_params::defrag_thold field and the --defrag-thold option.

Assisted-by: dsh:llama.cpp/ggml-org/DeepSeek-V4.1-Flash-1017-GGUF (Mac Studio, M5 Ultra, 256GB)

* fit-params : remove stale `llama-ext.h` include

Assisted-by: dsh:llama.cpp/ggml-org/DeepSeek-V4.1-Flash-1017-GGUF (Mac Studio, M5 Ultra, 256GB)

* llama : document `llama_model_dflash_selector_top_k`

Assisted-by: dsh:llama.cpp/ggml-org/DeepSeek-V4.1-Flash-1017-GGUF (Mac Studio, M5 Ultra, 256GB)

* llama : promote the NextN / speculative API from `llama-ext.h`

Move the NextN, layer input, embd row and state row APIs to `llama.h`
and rename `llama_set_embd_state_src` to `llama_set_ctx_state_src`, with
the internal `llama_context::set_state_src` following the same name.

Assisted-by: dsh:llama.cpp/ggml-org/DeepSeek-V4.1-Flash-1017-GGUF (Mac Studio, M5 Ultra, 256GB)

* llama : polish the NextN API promotion

Move the NextN draft model queries next to the Adapters section, rename
the `src` parameters of `llama_set_ctx_state_src` and
`llama_context::set_ctx_state_src` to `ctx_src`, and drop the now-redundant
`llama-ext.h` includes from `common/speculative.cpp` and `tools/imatrix`.

Assisted-by: dsh:llama.cpp/ggml-org/DeepSeek-V4.1-Flash-1017-GGUF (Mac Studio, M5 Ultra, 256GB)

.github/workflows/ui-build.yml

@@ -15,7 +15,7 @@ permissions:
 
 jobs:
   build:
-    runs-on: ubuntu-slim
+    runs-on: ubuntu-26.04-arm
 
     steps:
       - name: Checkout code

common/arg.cpp

@@ -943,8 +943,8 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context
     }
 
     const bool mcp_enabled = !params.mcp_servers_config.empty() || !params.mcp_servers_json.empty();
-    if ((!params.server_tools.empty() || mcp_enabled) && !params.cors_origins_explicit) {
-        LOG_WRN("server tools or MCP servers are enabled, using localhost as default CORS origin (change via --cors-origins)\n");
+    if ((!params.server_tools.empty() || mcp_enabled || params.ui_mcp_proxy) && !params.cors_origins_explicit) {
+        LOG_WRN("server tools, MCP servers or MCP proxy are enabled, using localhost as default CORS origin (change via --cors-origins)\n");
         params.cors_origins = "localhost";
     }
 
@@ -2501,15 +2501,6 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.ppl_output_type = value;
         }
     ).set_examples({LLAMA_EXAMPLE_PERPLEXITY}));
-    add_opt(common_arg(
-        {"-dt", "--defrag-thold"}, "N",
-        string_format("KV cache defragmentation threshold (DEPRECATED)"),
-        [](common_params & params, const std::string & value) {
-            GGML_UNUSED(params);
-            GGML_UNUSED(value);
-            LOG_WRN("DEPRECATED: --defrag-thold is deprecated and no longer necessary to specify\n");
-        }
-    ).set_env("LLAMA_ARG_DEFRAG_THOLD"));
     if (ex == LLAMA_EXAMPLE_SERVER) {
         // this is to make sure this option appears in the server-specific section of the help message
         add_opt(common_arg(
@@ -3361,7 +3352,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         {"--cors-origins"}, "ORIGINS",
         string_format(
             "comma-separated list of allowed origins for CORS (default: %s)\n"
-            "if set to special value 'localhost', reflect the Origin header only if it is localhost",
+            "if set to special value 'localhost', only allow same-origin requests (e.g. the built-in UI)",
         params.cors_origins.c_str()),
         [](common_params & params, const std::string & value) {
             params.cors_origins = value;
@@ -3417,7 +3408,8 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
     add_opt(common_arg(
         {"--ui-mcp-proxy", "--webui-mcp-proxy"},
         {"--no-ui-mcp-proxy", "--no-webui-mcp-proxy"},
-        "experimental: whether to enable MCP CORS proxy - do not enable in untrusted environments (default: disabled)",
+        "experimental: whether to enable MCP CORS proxy, for the Web UI to reach public MCP servers without CORS headers; local and private addresses are refused - do not enable in untrusted environments (default: disabled)\n"
+        "note: for security reasons, this will limit --cors-origins to localhost by default",
         [](common_params & params, bool value) {
             params.ui_mcp_proxy = value;
         }
@@ -3443,6 +3435,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.server_tools_runtime = value;
         }
     ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_TOOLS_RUNTIME"));
+    add_opt(common_arg(
+        {"--allow-unauthenticated-tools"},
+        "allow server tools, MCP servers or MCP proxy on a non-loopback address without an API key - only for trusted networks (default: disabled)",
+        [](common_params & params) {
+            params.server_tools_allow_unauth = true;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_ALLOW_UNAUTHENTICATED_TOOLS"));
     add_opt(common_arg(
         {"--mcp-servers-config"}, "PATH",
         "experimental: path to JSON file with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)\n"
@@ -3668,6 +3667,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             }
         }
     ).set_examples({LLAMA_EXAMPLE_SERVER}));
+    add_opt(common_arg(
+        {"--media-allow-private-urls"},
+        "allow http(s) media URLs (e.g. image_url) that point to local or private addresses (default: disabled, only public addresses)",
+        [](common_params & params) {
+            params.media_allow_private_urls = true;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MEDIA_ALLOW_PRIVATE_URLS"));
     add_opt(common_arg(
         {"--models-dir"}, "PATH",
         "directory containing models for the router server (default: disabled)",

common/chat-peg-parser.cpp

@@ -973,6 +973,21 @@ void common_chat_peg_gemma4_mapper::from_ast(const common_peg_ast_arena & arena,
     }
 }
 
+using gemma4 = common_chat_peg_gemma4_mapper;
+
+// Collects the nearest descendants carrying tag, looking through the rules the parser
+// wraps around the tagged nodes.
+static void gemma4_collect(const common_peg_ast_arena & arena, const common_peg_ast_node & node, const std::string & tag, std::vector<common_peg_ast_id> & out) {
+    for (auto child_id : node.children) {
+        const auto & child = arena.get(child_id);
+        if (child.tag == tag) {
+            out.push_back(child_id);
+        } else {
+            gemma4_collect(arena, child, tag, out);
+        }
+    }
+}
+
 static std::string gemma4_to_json(const common_peg_ast_arena & arena, common_peg_ast_id id) {
     const auto & node = arena.get(id);
 
@@ -980,51 +995,45 @@ static std::string gemma4_to_json(const common_peg_ast_arena & arena, common_peg
         return "";
     }
 
-    if (node.rule == "gemma4-number" || node.rule == "gemma4-bool" || node.rule == "gemma4-null") {
+    if (node.tag == gemma4::NUMBER || node.tag == gemma4::BOOL || node.tag == gemma4::NULL_) {
         return std::string(node.text);
     }
 
-    if (node.rule == "gemma4-string-content") {
+    if (node.tag == gemma4::STRING_CONTENT) {
         return escape_json_string_inner(node.sanitized_text());
     }
 
-    if (node.rule == "gemma4-string") {
+    if (node.tag == gemma4::STRING) {
         std::string result = "\"";
-        if (!node.children.empty()) {
-            result += gemma4_to_json(arena, node.children[0]);
-            if (!node.is_partial) {
-                result += "\"";
-            }
+        for (auto child_id : node.children) {
+            result += gemma4_to_json(arena, child_id);
+        }
+        if (!node.is_partial) {
+            result += "\"";
         }
         return result;
     }
 
-    if (node.rule == "gemma4-array") {
+    if (node.tag == gemma4::ARRAY) {
         std::string result = "[";
-
-        bool add_comma = false;
-        for (auto child_id : node.children) {
-            if (add_comma) {
+        for (size_t i = 0; i < node.children.size(); i++) {
+            if (i > 0) {
                 result += ',';
             }
-            add_comma = true;
-            result += gemma4_to_json(arena, child_id);
+            result += gemma4_to_json(arena, node.children[i]);
         }
-
         if (!node.is_partial) {
             result += ']';
         }
         return result;
     }
 
-    if (node.rule == "gemma4-dict-key-name") {
-        return std::string(node.text);
-    }
+    if (node.tag == gemma4::DICT_KEY) {
+        auto name_id = arena.find_by_tag(node, gemma4::DICT_KEY_NAME);
 
-    if (node.rule == "gemma4-dict-key") {
         std::string result = "\"";
-        if (!node.children.empty()) {
-            result += escape_json_string_inner(gemma4_to_json(arena, node.children[0]));
+        if (name_id != COMMON_PEG_INVALID_AST_ID) {
+            result += escape_json_string_inner(std::string(arena.get(name_id).text));
         }
         if (!node.is_partial) {
             result += "\":";
@@ -1032,40 +1041,29 @@ static std::string gemma4_to_json(const common_peg_ast_arena & arena, common_peg
         return result;
     }
 
-    if (node.rule == "gemma4-dict-kv") {
-        std::string result;
-        for (auto child_id : node.children) {
-            result += gemma4_to_json(arena, child_id);
-        }
-        return result;
-    }
+    if (node.tag == gemma4::DICT) {
+        std::vector<common_peg_ast_id> members;
+        gemma4_collect(arena, node, gemma4::DICT_KV, members);
 
-    if (node.rule == "gemma4-dict") {
         std::string result = "{";
-
-        bool add_comma = false;
-        for (auto child_id : node.children) {
-            if (add_comma) {
+        for (size_t i = 0; i < members.size(); i++) {
+            if (i > 0) {
                 result += ',';
             }
-            add_comma = true;
-            result += gemma4_to_json(arena, child_id);
+            result += gemma4_to_json(arena, members[i]);
         }
-
         if (!node.is_partial) {
             result += '}';
         }
         return result;
     }
 
-    if (node.rule == "gemma4-value") {
-        if (!node.children.empty()) {
-            return gemma4_to_json(arena, node.children[0]);
-        }
-        return "";
+    // members, rules and other wrappers
+    std::string result;
+    for (auto child_id : node.children) {
+        result += gemma4_to_json(arena, child_id);
     }
-
-    return "";
+    return result;
 }
 
 void common_chat_peg_gemma4_mapper::visit(const common_peg_ast_arena & arena, common_peg_ast_id id) {

common/chat-peg-parser.h

@@ -34,6 +34,18 @@ class common_chat_peg_mapper {
 
 class common_chat_peg_gemma4_mapper : public common_chat_peg_mapper {
   public:
+    // Tags of the Gemma 4 argument syntax. The parser may wrap them in any rules.
+    static constexpr const char * STRING         = "gemma4-string";
+    static constexpr const char * STRING_CONTENT = "gemma4-string-content";
+    static constexpr const char * NUMBER         = "gemma4-number";
+    static constexpr const char * BOOL           = "gemma4-bool";
+    static constexpr const char * NULL_          = "gemma4-null";
+    static constexpr const char * ARRAY          = "gemma4-array";
+    static constexpr const char * DICT           = "gemma4-dict";
+    static constexpr const char * DICT_KV        = "gemma4-dict-kv";
+    static constexpr const char * DICT_KEY       = "gemma4-dict-key";
+    static constexpr const char * DICT_KEY_NAME  = "gemma4-dict-key-name";
+
     common_chat_peg_gemma4_mapper(common_chat_msg & msg) : common_chat_peg_mapper(msg) {}
     virtual void from_ast(const common_peg_ast_arena & arena, const common_peg_parse_result & result);
   private:

common/chat.cpp

@@ -1277,13 +1277,6 @@ static common_chat_params common_chat_params_init_lfm2_5(const common_chat_templ
     return common_chat_params_init_lfm2(tmpl, inputs, /* tool_list_tokens = */ false);
 }
 
-// Older gemma4 templates need their tool responses rewritten before rendering
-static common_chat_params common_chat_params_init_gemma4_legacy(const common_chat_template & tmpl, const autoparser::generation_params & inputs) {
-    auto adjusted = inputs;
-    workaround::convert_tool_responses_gemma4(adjusted.messages);
-    return common_chat_params_init_gemma4(tmpl, adjusted);
-}
-
 // Pick the dedicated handler for a template from its source, or null for the autoparser.
 // Order matters: the first match wins, and later checks assume the earlier ones did not match.
 static common_chat_params_init_fn common_chat_template_detect_params_init(const std::string & src) {
@@ -1418,11 +1411,6 @@ static common_chat_params_init_fn common_chat_template_detect_params_init(const
     // Gemma4 format detection
     if (src.find("'<|tool_call>call:'") != std::string::npos) {
         LOG_DBG("Using specialized template: Gemma4\n");
-        if (src.find("{#- OpenAI Chat Completions:") == std::string::npos) {
-            LOG_WRN("%s: detected an outdated gemma4 chat template, applying compatibility workarounds. "
-                    "Consider updating to the official template.\n", __func__);
-            return common_chat_params_init_gemma4_legacy;
-        }
         return common_chat_params_init_gemma4;
     }
 

common/common.h

@@ -705,6 +705,7 @@ struct common_params {
     // enable built-in tools
     std::vector<std::string> server_tools;
     std::string server_tools_runtime;
+    bool server_tools_allow_unauth = false; // allow tools/MCP on a non-loopback address without API key
 
     // MCP server configs (Cursor-compatible JSON)
     std::string mcp_servers_config;   // path to JSON file with MCP server definitions
@@ -722,6 +723,7 @@ struct common_params {
 
     std::string slot_save_path;
     std::string media_path; // path to directory for loading media files
+    bool media_allow_private_urls = false; // allow http(s) media URLs on local or private addresses
 
     float slot_prompt_similarity = 0.1f;
 
@@ -1147,7 +1149,7 @@ struct common_batch {
     // attach a state embedding (e.g. the target hidden state for MTP) to the entry at idx, can only be set once per entry
     bool set_embd_state(int32_t idx, llama_embd state);
 
-    // the state of entry idx is row `row` of the h_nextn device buffer of the state source context, see llama_set_embd_state_src()
+    // the state of entry idx is row `row` of the h_nextn device buffer of the state source context, see llama_set_ctx_state_src()
     bool set_embd_state_row(int32_t idx, int32_t row);
 
     // add an embedding-only entry (no token id)

common/download.cpp

@@ -461,38 +461,55 @@ static int common_download_file_single_online(const std::string & url,
     return finish(success ? head->status : -1);
 }
 
-std::pair<long, std::vector<char>> common_remote_get_content(const std::string          & url,
+std::pair<long, std::vector<char>> common_remote_get_content(const std::string          & url_in,
                                                              const common_remote_params & params) {
-    auto [cli, parts] = common_http_client(url);
+    static constexpr int max_redirects = 5;
+    std::string url = url_in;
+    for (int n_redirects = 0; ; n_redirects++) {
+        auto [cli, parts] = common_http_client(url);
 
-    httplib::Headers headers;
-    for (const auto & h : params.headers) {
-        headers.emplace(h.first, h.second);
+        httplib::Headers headers;
+        for (const auto & h : params.headers) {
+            headers.emplace(h.first, h.second);
+        }
+        if (headers.find("User-Agent") == headers.end()) {
+            headers.emplace("User-Agent", "llama-cpp/" + std::string(llama_build_info()));
+        }
+
+        if (params.timeout > 0) {
+            cli.set_read_timeout(params.timeout, 0);
+            cli.set_write_timeout(params.timeout, 0);
+        }
+        if (params.resolve_host) {
+            // pin checked address, follow redirects manually to check each hop
+            cli.set_direct_hostname_addr_map({{parts.host, params.resolve_host(parts.host)}});
+            cli.set_follow_location(false);
+        }
+
+        std::vector<char> buf;
+        auto res = cli.Get(parts.path, headers,
+            [&](const char *data, size_t len) {
+                buf.insert(buf.end(), data, data + len);
+                return params.max_size == 0 ||
+                       buf.size() <= static_cast<size_t>(params.max_size);
+            },
+            nullptr
+        );
+
+        if (!res) {
+            throw std::runtime_error("error: cannot make GET request");
+        }
+
+        if (params.resolve_host && res->status >= 300 && res->status < 400 && res->has_header("Location")) {
+            if (n_redirects >= max_redirects) {
+                throw std::runtime_error("error: too many redirects");
+            }
+            url = common_http_resolve_location(url, res->get_header_value("Location"));
+            continue;
+        }
+
+        return { res->status, std::move(buf) };
     }
-    if (headers.find("User-Agent") == headers.end()) {
-        headers.emplace("User-Agent", "llama-cpp/" + std::string(llama_build_info()));
-    }
-
-    if (params.timeout > 0) {
-        cli.set_read_timeout(params.timeout, 0);
-        cli.set_write_timeout(params.timeout, 0);
-    }
-
-    std::vector<char> buf;
-    auto res = cli.Get(parts.path, headers,
-        [&](const char *data, size_t len) {
-            buf.insert(buf.end(), data, data + len);
-            return params.max_size == 0 ||
-                   buf.size() <= static_cast<size_t>(params.max_size);
-        },
-        nullptr
-    );
-
-    if (!res) {
-        throw std::runtime_error("error: cannot make GET request");
-    }
-
-    return { res->status, std::move(buf) };
 }
 
 int common_download_file_single(const std::string & url,

common/download.h

@@ -31,6 +31,8 @@ struct common_remote_params {
     common_header_list headers;
     long timeout  = 0;           // in seconds, 0 means no timeout
     long max_size = 0;           // unlimited if 0
+    // if set: called per hop (redirects too), returns the address to connect to, throws to refuse
+    std::function<std::string(const std::string & host)> resolve_host;
 };
 
 // get remote file content, returns <http_code, raw_response_body>

common/http.h

@@ -99,6 +99,26 @@ static common_http_url common_http_parse_url(const std::string & url) {
     return parts;
 }
 
+// resolve a redirect Location against the URL it came from
+static std::string common_http_resolve_location(const std::string & base_url, const std::string & location) {
+    if (location.find("://") != std::string::npos) {
+        return location;
+    }
+    const common_http_url base = common_http_parse_url(base_url);
+    if (location.compare(0, 2, "//") == 0) {
+        return base.scheme + ":" + location;
+    }
+    const std::string origin = base.scheme + "://" + common_http_format_host(base.host) + ":" + std::to_string(base.port);
+    const std::string path = base.path.substr(0, base.path.find('?'));
+    if (!location.empty() && location[0] == '/') {
+        return origin + location;
+    }
+    if (!location.empty() && location[0] == '?') {
+        return origin + path + location;
+    }
+    return origin + path.substr(0, path.rfind('/') + 1) + location;
+}
+
 // match a loopback destination: a remote proxy can never reach the client machine's own
 // loopback, so these hosts are always contacted directly, before NO_PROXY is consulted
 // (similar to Go's net/http)

common/json-schema.h

@@ -85,92 +85,124 @@ struct common_chat_schema {
 
     static const char * kind_name(node_kind kind);
     static const char * type_name(value_type type);
+
+    template <typename T>
+    const T & as() const {
+        GGML_ASSERT(kind() == T::KIND);
+        return static_cast<const T &>(*this);
+    }
 };
 
 using common_chat_schema_ptr = std::unique_ptr<common_chat_schema>;
 
 struct common_chat_schema_any : common_chat_schema {
-    node_kind kind() const override { return KIND_ANY; }
+    static constexpr node_kind KIND = KIND_ANY;
+
+    node_kind kind() const override { return KIND; }
 };
 
 // {"$ref": "#/..."}, only references into the same document are supported
 struct common_chat_schema_ref : common_chat_schema {
+    static constexpr node_kind KIND = KIND_REF;
+
     std::string                ref;
     const common_chat_schema * target = nullptr;  // owned by common_chat_schema_document::refs
 
     explicit common_chat_schema_ref(std::string ref) : ref(std::move(ref)) {}
 
-    node_kind kind() const override { return KIND_REF; }
+    node_kind kind() const override { return KIND; }
 };
 
 // oneOf / anyOf, or a "type" array expanded to one alternative per type
 struct common_chat_schema_any_of : common_chat_schema {
+    static constexpr node_kind KIND = KIND_ANY_OF;
+
     std::vector<common_chat_schema_ptr> children;
 
-    node_kind kind() const override { return KIND_ANY_OF; }
+    node_kind kind() const override { return KIND; }
 };
 
 struct common_chat_schema_all_of : common_chat_schema {
+    static constexpr node_kind KIND = KIND_ALL_OF;
+
     std::vector<common_chat_schema_ptr> children;
 
-    node_kind kind() const override { return KIND_ALL_OF; }
+    node_kind kind() const override { return KIND; }
 };
 
 struct common_chat_schema_const : common_chat_schema {
+    static constexpr node_kind KIND = KIND_CONST;
+
     common_json value;
 
     explicit common_chat_schema_const(common_json value) : value(std::move(value)) {}
 
-    node_kind kind() const override { return KIND_CONST; }
+    node_kind kind() const override { return KIND; }
 };
 
 struct common_chat_schema_enum : common_chat_schema {
+    static constexpr node_kind KIND = KIND_ENUM;
+
     std::vector<common_json> values;
 
-    node_kind kind() const override { return KIND_ENUM; }
+    node_kind kind() const override { return KIND; }
 };
 
 struct common_chat_schema_null : common_chat_schema {
-    node_kind kind() const override { return KIND_NULL; }
+    static constexpr node_kind KIND = KIND_NULL;
+
+    node_kind kind() const override { return KIND; }
 };
 
 struct common_chat_schema_boolean : common_chat_schema {
-    node_kind kind() const override { return KIND_BOOLEAN; }
+    static constexpr node_kind KIND = KIND_BOOLEAN;
+
+    node_kind kind() const override { return KIND; }
 };
 
 struct common_chat_schema_number : common_chat_schema {
-    node_kind kind() const override { return KIND_NUMBER; }
+    static constexpr node_kind KIND = KIND_NUMBER;
+
+    node_kind kind() const override { return KIND; }
 };
 
 // bounds are inclusive, exclusiveMinimum / exclusiveMaximum are folded in
 struct common_chat_schema_integer : common_chat_schema {
+    static constexpr node_kind KIND = KIND_INTEGER;
+
     int64_t minimum = INT64_MIN;  // INT64_MIN for unbounded
     int64_t maximum = INT64_MAX;  // INT64_MAX for unbounded
 
-    node_kind kind() const override { return KIND_INTEGER; }
+    node_kind kind() const override { return KIND; }
 };
 
 struct common_chat_schema_string : common_chat_schema {
+    static constexpr node_kind KIND = KIND_STRING;
+
     std::string   pattern;  // empty when absent
     string_format format     = FORMAT_NONE;
     int           min_length = 0;
     int           max_length = -1;  // -1 for unbounded
 
-    node_kind kind() const override { return KIND_STRING; }
+    node_kind kind() const override { return KIND; }
 };
 
 struct common_chat_schema_array : common_chat_schema {
+    static constexpr node_kind KIND = KIND_ARRAY;
+
     common_chat_schema_ptr items;  // a common_chat_schema_any when "items" is absent
     int                    min_items = 0;
     int                    max_items = -1;  // -1 for unbounded
 
-    node_kind kind() const override { return KIND_ARRAY; }
+    node_kind kind() const override { return KIND; }
 };
 
 struct common_chat_schema_tuple : common_chat_schema {
+    static constexpr node_kind KIND = KIND_TUPLE;
+
     std::vector<common_chat_schema_ptr> items;
 
-    node_kind kind() const override { return KIND_TUPLE; }
+    node_kind kind() const override { return KIND; }
 };
 
 struct common_chat_schema_property {
@@ -180,10 +212,12 @@ struct common_chat_schema_property {
 };
 
 struct common_chat_schema_object : common_chat_schema {
+    static constexpr node_kind KIND = KIND_OBJECT;
+
     std::vector<common_chat_schema_property> properties;             // in schema order
     common_chat_schema_ptr                   additional_properties;  // null when not allowed
 
-    node_kind kind() const override { return KIND_OBJECT; }
+    node_kind kind() const override { return KIND; }
 };
 
 struct common_chat_schema_document {

common/parsers/gemma4.cpp

@@ -1,152 +1,284 @@
 #include "parsers.h"
 
-namespace workaround {
+#include "log.h"
 
-// Gemma4 uses a custom tool_responses field instead of role:tool messages.
-//
-// This will transform a sequence of messages:
-//   assistant(tool_call+) -> tool+ -> assistant(content)
-//
-// Into a single assistant message containing a tool_responses field:
-//   assistant(content + tool_call + tool_responses)
-//
-// This is necessary for the Gemma4 chat template to properly format the prompt.
-// See https://ai.google.dev/gemma/docs/core/prompt-formatting-gemma4
-struct gemma4_model_turn_builder {
-    json & messages;
-    size_t pos;
-    json tool_calls = json::array();
-    json tool_responses = json::array();
-    json content;
-    json reasoning_content;
+#include <algorithm>
+#include <cctype>
+#include <map>
+#include <optional>
 
-    gemma4_model_turn_builder(json & msgs, size_t pos) : messages(msgs), pos(pos) {}
+namespace {
 
-    void collect() {
-        // Collect the first assistant message
-        auto & msg = messages[pos];
-        if (msg.contains("reasoning_content") && msg.at("reasoning_content").is_string()) {
-            // According to the prompt formatting guide, we need to preserve reasoning_content
-            // between function calls. The current chat templates do not support this, but we will do it anyway.
-            reasoning_content = msg.at("reasoning_content");
+using gemma4 = common_chat_peg_gemma4_mapper;
+
+struct depth_guard {
+    int & depth;
+    depth_guard(int & depth) : depth(depth) {
+        if (++depth > COMMON_RECURSIVE_LIMIT) {
+            throw std::invalid_argument("JSON schema is nested more than " + std::to_string(COMMON_RECURSIVE_LIMIT) + " levels deep");
         }
-        for (auto & tc : msg.at("tool_calls")) {
-            tool_calls.push_back(tc);
-        }
-        pos++;
+    }
+    ~depth_guard() { depth--; }
+};
 
-        // Collect tool call results
-        while (pos < messages.size() && messages[pos].value("role", "") == "tool") {
-            collect_result(messages[pos]);
-            pos++;
-        }
+// Gemma 4 value parsers
+void gemma4_value_rules(common_chat_peg_builder & p) {
+    // The tags are what gemma4_to_json() in the mapper reads, the rules may wrap them freely.
+    p.rule("gemma4-string", p.tag(gemma4::STRING, p.literal("<|\"|>") + p.tag(gemma4::STRING_CONTENT, p.until("<|\"|>")) + p.literal("<|\"|>")));
 
-        // Check if the next assistant message is the final message
-        if (pos < messages.size() && messages[pos].value("role", "") == "assistant") {
-            auto & next = messages[pos];
-            if (!has_tool_calls(next) && has_content(next)) {
-                content = next.at("content");
-                pos++;
-            }
+    p.rule("gemma4-bool", p.tag(gemma4::BOOL, p.json_bool()));
+
+    p.rule("gemma4-null", p.tag(gemma4::NULL_, p.json_null()));
+
+    p.rule("gemma4-number", p.tag(gemma4::NUMBER, p.json_number()));
+
+    // Keys are bare in the template, but Gemma 4 occasionally wraps them in quotes
+    p.rule("gemma4-dict-key", p.tag(gemma4::DICT_KEY, p.choice({
+        p.literal("<|\"|>") + p.tag(gemma4::DICT_KEY_NAME, p.until("<|\"|>")) + p.literal("<|\"|>"),
+        p.literal("\"") + p.tag(gemma4::DICT_KEY_NAME, p.chars("[^\"]", 1, -1)) + p.literal("\""),
+        p.literal("'") + p.tag(gemma4::DICT_KEY_NAME, p.chars("[^']", 1, -1)) + p.literal("'"),
+        p.tag(gemma4::DICT_KEY_NAME, p.chars("[^:}]", 1, -1)),
+    }) + p.literal(":")));
+
+    p.rule("gemma4-dict-kv", p.tag(gemma4::DICT_KV, p.ref("gemma4-dict-key") + p.space() + p.ref("gemma4-value")));
+
+    p.rule("gemma4-dict", [&]() {
+        auto ws = p.space();
+        auto member = p.ref("gemma4-dict-kv");
+        auto members = p.sequence({member, p.zero_or_more(p.sequence({p.literal(","), ws, member}))});
+        return p.tag(gemma4::DICT, p.sequence({
+            p.literal("{"), ws,
+            p.choice({p.literal("}"), p.sequence({members, ws, p.literal("}")})})
+        }));
+    });
+
+    p.rule("gemma4-array", [&]() {
+        auto ws = p.space();
+        auto value = p.ref("gemma4-value");
+        auto elements = p.sequence({value, p.zero_or_more(p.sequence({p.literal(","), ws, value}))});
+        return p.tag(gemma4::ARRAY, p.sequence({
+            p.literal("["), ws,
+            p.choice({p.literal("]"), p.sequence({elements, ws, p.literal("]")})})
+        }));
+    });
+
+    p.rule("gemma4-value", [&]() {
+        return p.choice({
+            p.ref("gemma4-string"), p.ref("gemma4-dict"), p.ref("gemma4-array"),
+            p.ref("gemma4-number"), p.ref("gemma4-bool"), p.ref("gemma4-null")
+        });
+    });
+}
+
+// Gemma 4 uses different delimiters and cannot leverage the json-schema-to-grammar implementation.
+// This implements a subset to enforce properties.
+class gemma4_schema_parser {
+    common_chat_peg_builder & p_;
+    std::string               root_;
+    int                       depth_ = 0;
+
+    std::map<const common_chat_schema *, std::string> refs_;
+
+  public:
+    gemma4_schema_parser(common_chat_peg_builder & p, std::string root) : p_(p), root_(std::move(root)) {}
+
+    common_peg_parser visit(const common_chat_schema & node) { return visit(node, root_); }
+
+  private:
+    common_peg_parser visit(const common_chat_schema & node, const std::string & name) {
+        depth_guard guard(depth_);
+        switch (node.kind()) {
+            case common_chat_schema::KIND_REF:    return visit_ref(node.as<common_chat_schema_ref>());
+            case common_chat_schema::KIND_ANY_OF: return visit_any_of(node.as<common_chat_schema_any_of>(), name);
+            case common_chat_schema::KIND_ALL_OF: return visit_all_of(node.as<common_chat_schema_all_of>(), name);
+            case common_chat_schema::KIND_ARRAY:  return visit_array(node.as<common_chat_schema_array>(), name);
+            case common_chat_schema::KIND_TUPLE:  return visit_tuple(node.as<common_chat_schema_tuple>(), name);
+            case common_chat_schema::KIND_OBJECT: return visit_object(node.as<common_chat_schema_object>(), name);
+            default:                              return p_.ref("gemma4-value");
         }
     }
 
-    void collect_result(const json & curr) {
-        json response;
-        if (curr.contains("content")) {
-            const auto & content = curr.at("content");
-            if (content.is_string()) {
-                // Try to parse the content as JSON; fall back to raw string
-                try {
-                    response = json::parse(content.get<std::string>());
-                } catch (...) {
-                    response = content;
+    // Whether a schema contributes a shape of its own, as opposed to accepting any value.
+    static bool shaped(const common_chat_schema & node) {
+        switch (node.kind()) {
+            case common_chat_schema::KIND_REF:
+                return node.as<common_chat_schema_ref>().target != nullptr;
+            case common_chat_schema::KIND_ANY_OF:
+                for (const auto & child : node.as<common_chat_schema_any_of>().children) {
+                    if (shaped(*child)) {
+                        return true;
+                    }
                 }
+                return false;
+            case common_chat_schema::KIND_ALL_OF: {
+                const auto & children = node.as<common_chat_schema_all_of>().children;
+                return children.size() == 1 && shaped(*children[0]);
+            }
+            case common_chat_schema::KIND_ARRAY:
+            case common_chat_schema::KIND_TUPLE:
+            case common_chat_schema::KIND_OBJECT:
+                return true;
+            default:
+                return false;
+        }
+    }
+
+    // The template writes keys bare. Models keep that for identifiers and quote anything else,
+    // with single or double quotes or the string delimiter, so those are accepted as well.
+    common_peg_parser key(const std::string & key) {
+        auto name = [&]() { return p_.tag(gemma4::DICT_KEY_NAME, p_.literal(key)); };
+
+        bool identifier = !key.empty() && !std::isdigit(static_cast<unsigned char>(key[0]));
+        for (char c : key) {
+            if (!std::isalnum(static_cast<unsigned char>(c)) && c != '_') {
+                identifier = false;
+            }
+        }
+        if (identifier) {
+            return p_.tag(gemma4::DICT_KEY, name() + p_.literal(":"));
+        }
+
+        auto alts = p_.choice();
+        if (!key.empty() && key.find_first_of(":}") == std::string::npos && std::string("'\"<").find(key[0]) == std::string::npos) {
+            // Check if key is unambiguous
+            alts |= name();
+        }
+        // Gemma 4 may emit keys in its string delimiters or quotes
+        alts |= p_.literal("<|\"|>") + name() + p_.literal("<|\"|>");
+        if (key.find('\'') == std::string::npos) {
+            alts |= p_.literal("'") + name() + p_.literal("'");
+        }
+        if (key.find('"') == std::string::npos) {
+            alts |= p_.literal("\"") + name() + p_.literal("\"");
+        }
+        return p_.tag(gemma4::DICT_KEY, alts + p_.literal(":"));
+    }
+
+    common_peg_parser visit_ref(const common_chat_schema_ref & ref) {
+        if (!ref.target) {
+            return p_.ref("gemma4-value");
+        }
+        auto it = refs_.find(ref.target);
+        if (it != refs_.end()) {
+            return p_.ref(it->second);
+        }
+        // A target reached from several places gets one rule, named by the order it was first seen.
+        auto rule = root_ + "-ref-" + std::to_string(refs_.size());
+        refs_[ref.target] = rule;
+        return p_.rule(rule, [&]() { return visit(*ref.target, rule + "-value"); });
+    }
+
+    // Shaped alternatives are tried first, then anything else collapses into one generic value.
+    common_peg_parser visit_any_of(const common_chat_schema_any_of & node, const std::string & name) {
+        if (!shaped(node)) {
+            return p_.ref("gemma4-value");
+        }
+        if (node.children.size() == 1) {
+            return visit(*node.children[0], name);
+        }
+        auto alts  = p_.choice();
+        bool loose = false;
+        for (size_t i = 0; i < node.children.size(); i++) {
+            if (shaped(*node.children[i])) {
+                alts |= visit(*node.children[i], name + "-" + std::to_string(i));
             } else {
-                response = content;
+                loose = true;
             }
         }
-
-        std::string name;
-
-        // Match name with corresponding tool call
-        size_t idx = tool_responses.size();
-        if (idx < tool_calls.size()) {
-            auto & tc = tool_calls[idx];
-            if (tc.contains("function")) {
-                name = tc.at("function").value("name", "");
-            }
+        if (loose) {
+            alts |= p_.ref("gemma4-value");
         }
-
-        // Fallback to the tool call id
-        if (name.empty()) {
-            name = curr.value("tool_call_id", "");
-        }
-
-        tool_responses.push_back({{"name", name}, {"response", response}});
+        return p_.rule(name, alts);
     }
 
-    json build() {
-        collect();
+    common_peg_parser visit_all_of(const common_chat_schema_all_of & node, const std::string & name) {
+        if (node.children.size() == 1) {
+            return visit(*node.children[0], name);
+        }
+        return p_.ref("gemma4-value");
+    }
 
-        json msg = {
-            {"role", "assistant"},
-            {"tool_calls", tool_calls},
+    common_peg_parser visit_array(const common_chat_schema_array & node, const std::string & name) {
+        auto ws       = p_.space();
+        auto item     = node.items && shaped(*node.items) ? visit(*node.items, name + "-item") : p_.ref("gemma4-value");
+        auto elements = item + p_.zero_or_more(p_.literal(",") + ws + item);
+        return p_.rule(name, p_.tag(gemma4::ARRAY,
+            p_.literal("[") + ws + (p_.literal("]") | (elements + ws + p_.literal("]")))));
+    }
+
+    common_peg_parser visit_tuple(const common_chat_schema_tuple & node, const std::string & name) {
+        auto ws    = p_.space();
+        auto items = p_.sequence();
+        for (size_t i = 0; i < node.items.size(); i++) {
+            if (i > 0) {
+                items = items + p_.literal(",") + ws;
+            }
+            items = items + visit(*node.items[i], name + "-" + std::to_string(i));
+        }
+        return p_.rule(name, p_.tag(gemma4::ARRAY, p_.literal("[") + ws + items + ws + p_.literal("]")));
+    }
+
+    common_peg_parser visit_object(const common_chat_schema_object & node, const std::string & name) {
+        auto ws = p_.space();
+
+        std::vector<const common_chat_schema_property *> props;
+        for (const auto & prop : node.properties) {
+            props.push_back(&prop);
+        }
+
+        // The Gemma 4 jinja template sorts properties, so our rules need to match.
+        auto lower = [](const std::string & s) {
+            std::string out = s;
+            for (auto & c : out) {
+                c = static_cast<char>(std::tolower(static_cast<unsigned char>(c)));
+            }
+            return out;
         };
-        if (!tool_responses.empty()) {
-            msg["tool_responses"] = tool_responses;
-        }
-        if (!content.is_null()) {
-            msg["content"] = content;
-        }
-        if (!reasoning_content.is_null()) {
-            msg["reasoning_content"] = reasoning_content;
-        }
-        return msg;
-    }
+        std::stable_sort(props.begin(), props.end(), [&](const auto * a, const auto * b) {
+            return lower(a->name) < lower(b->name);
+        });
 
-    static bool has_content(const json & msg) {
-        if (!msg.contains("content") || msg.at("content").is_null()) {
-            return false;
+        if (props.empty()) {
+            if (node.additional_properties) {
+                return p_.ref("gemma4-dict");
+            }
+            return p_.rule(name, p_.tag(gemma4::DICT, p_.literal("{") + ws + p_.literal("}")));
         }
-        const auto & content = msg.at("content");
-        if (content.is_string() && !content.get<std::string>().empty()) {
-            return true;
-        }
-        if (content.is_array() && !content.empty()) {
-            return true;
-        }
-        return false;
-    }
 
-    static bool has_tool_calls(const json & msg) {
-        return msg.contains("tool_calls") && msg.at("tool_calls").is_array() && !msg.at("tool_calls").empty();
+        std::optional<common_peg_parser> rest;
+        std::optional<common_peg_parser> first;
+        if (node.additional_properties) {
+            auto kv = p_.ref("gemma4-dict-kv");
+            rest    = p_.zero_or_more(p_.literal(",") + ws + kv);
+            first   = p_.optional(kv + *rest);
+        }
+
+        auto then_rest = [&](const common_peg_parser & parser) {
+            return rest ? parser + *rest : parser;
+        };
+
+        for (size_t i = props.size(); i-- > 0;) {
+            const auto & prop = *props[i];
+            auto kv = p_.tag(gemma4::DICT_KV, key(prop.name) + ws + visit(*prop.schema, name + "-" + std::to_string(i)));
+
+            common_peg_parser rest_i  = then_rest(p_.literal(",") + ws + kv);
+            common_peg_parser first_i = then_rest(kv);
+            if (!prop.required) {
+                rest_i  = then_rest(p_.optional(p_.literal(",") + ws + kv));
+                first_i = first ? (first_i | *first) : p_.optional(first_i);
+            }
+
+            if (i > 0) {
+                rest = p_.rule(name + "-rest-" + std::to_string(i), rest_i);
+            }
+            first = first_i;
+        }
+
+        return p_.rule(name, p_.tag(gemma4::DICT, p_.literal("{") + ws + *first + ws + p_.literal("}")));
     }
 };
 
-void convert_tool_responses_gemma4(json & messages) {
-    json result = json::array();
-    size_t i = 0;
-
-    while (i < messages.size()) {
-        auto & msg = messages[i];
-
-        if (msg.value("role", "") != "assistant" || !msg.contains("tool_calls") ||
-            !msg.at("tool_calls").is_array() || msg.at("tool_calls").empty()) {
-            result.push_back(msg);
-            i++;
-            continue;
-        }
-
-        gemma4_model_turn_builder builder(messages, i);
-        result.push_back(builder.build());
-        i = builder.pos;
-    }
-
-    messages = result;
-}
-
-}
+}  // namespace
 
 common_chat_params common_chat_params_init_gemma4(const common_chat_template &    tmpl,
                                                          const autoparser::generation_params & inputs) {
@@ -218,51 +350,29 @@ common_chat_params common_chat_params_init_gemma4(const common_chat_template &
         }
 
         if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
-            // Gemma4 tool calling syntax
-            // Rules should match traversal logic in gemma4_to_json()
-            p.rule("gemma4-string-content", p.until("<|\"|>"));
-            p.rule("gemma4-string", p.literal("<|\"|>") + p.ref("gemma4-string-content") + p.literal("<|\"|>"));
-            p.rule("gemma4-bool", p.json_bool());
-            p.rule("gemma4-null", p.json_null());
-            p.rule("gemma4-number", p.json_number());
-            p.rule("gemma4-dict-key", p.rule("gemma4-dict-key-name", p.chars("[^:}]", 1, -1)) + p.literal(":"));
-            p.rule("gemma4-dict-kv", p.ref("gemma4-dict-key") + p.space() + p.ref("gemma4-value"));
-            p.rule("gemma4-dict", [&]() {
-                auto ws = p.space();
-                auto member = p.ref("gemma4-dict-kv");
-                auto members = p.sequence({member, p.zero_or_more(p.sequence({p.literal(","), ws, member}))});
-                return p.sequence({
-                    p.literal("{"), ws,
-                    p.choice({p.literal("}"), p.sequence({members, ws, p.literal("}")})})
-                });
-            });
-            p.rule("gemma4-array", [&]() {
-                auto ws = p.space();
-                auto value = p.ref("gemma4-value");
-                auto elements = p.sequence({value, p.zero_or_more(p.sequence({p.literal(","), ws, value}))});
-                return p.sequence({
-                    p.literal("["), ws,
-                    p.choice({p.literal("]"), p.sequence({elements, ws, p.literal("]")})})
-                });
-            });
-            p.rule("gemma4-value", [&]() {
-                return p.choice({
-                    p.ref("gemma4-string"), p.ref("gemma4-dict"), p.ref("gemma4-array"),
-                    p.ref("gemma4-number"), p.ref("gemma4-bool"), p.ref("gemma4-null")
-                });
-            });
+            gemma4_value_rules(p);
 
             auto tool_choice = p.choice();
 
             foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                 const auto & function = tool.at("function");
                 std::string  name     = function.at("name");
-                // TODO @aldehir : need to extend json-schema-to-grammar to produce more than JSON rules
-                // const auto & params   = function.at("parameters");
+
+                // Arguments follow the shape of the parameter schema when it is an object,
+                // see gemma4_schema_parser, and are a generic dict otherwise.
+                auto args = p.ref("gemma4-dict");
+                try {
+                    auto doc = common_chat_schema_from_json(common_chat_tool_parameters(function));
+                    if (doc.root->kind() == common_chat_schema::KIND_OBJECT) {
+                        args = gemma4_schema_parser(p, "tool-" + std::to_string(tool_index) + "-args").visit(*doc.root);
+                    }
+                } catch (const std::exception & e) {
+                    LOG_WRN("%s: tool %s uses the generic argument parser: %s\n", __func__, name.c_str(), e.what());
+                }
 
                 tool_choice |= p.rule("tool-" + std::to_string(tool_index), p.tool(p.sequence({
                     p.tool_open(p.tool_name(p.literal(name)) + p.peek(p.literal("{"))),
-                    p.tool_args(p.ref("gemma4-dict")),
+                    p.tool_args(args),
                 })));
             });
 

common/parsers/parsers.h

@@ -41,12 +41,6 @@ std::string common_chat_template_generation_prompt_impl(
 
 bool is_lfm2_template(const std::string & src);
 
-namespace workaround {
-
-void convert_tool_responses_gemma4(json & messages);
-
-}
-
 common_chat_params common_chat_params_init_cohere2moe(const common_chat_template & tmpl, const autoparser::generation_params & inputs);
 
 common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_template & tmpl, const autoparser::generation_params & inputs);

common/speculative.cpp

@@ -10,8 +10,6 @@
 #include "ngram-mod.h"
 #include "sampling.h"
 
-#include "../src/llama-ext.h" // staging API: llama_set_embeddings_nextn / llama_get_embeddings_nextn_ith (used by MTP)
-
 #include <algorithm>
 #include <cassert>
 #include <cmath>
@@ -1650,9 +1648,9 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
 
         // shared memory and chained heads keep the host rows
         if (!is_mem_shared && !chain_heads) {
-            dev_rows = llama_set_embeddings_nextn_dev(ctx_tgt, true) && llama_set_embd_state_src(ctx_dft, ctx_tgt);
+            dev_rows = llama_set_embeddings_nextn_dev(ctx_tgt, true) && llama_set_ctx_state_src(ctx_dft, ctx_tgt);
             if (!dev_rows) {
-                llama_set_embd_state_src(ctx_dft, nullptr);
+                llama_set_ctx_state_src(ctx_dft, nullptr);
                 llama_set_embeddings_nextn_dev(ctx_tgt, false);
             }
         }
@@ -1662,7 +1660,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
     ~common_speculative_impl_draft_mtp() override {
         auto * ctx_dft = this->params.ctx_dft;
         if (dev_rows) {
-            llama_set_embd_state_src(ctx_dft, nullptr);
+            llama_set_ctx_state_src(ctx_dft, nullptr);
             llama_set_embeddings_nextn_dev(this->params.ctx_tgt, false);
         }
         for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) backend_chains.size(); ++seq_id) {

conversion/qwen.py

@@ -493,6 +493,45 @@ class _LinearAttentionVReorderBase(Qwen3NextModel):
     interleaved repeat: [G0_v0, G1_v0, ..., G0_v1, G1_v1, ...].
     """
 
+    _SHARED_INPUT_MERGES: tuple[tuple[gguf.MODEL_TENSOR, tuple[gguf.MODEL_TENSOR, gguf.MODEL_TENSOR], bool], ...] = ()
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._merge_parts: dict[tuple[int, gguf.MODEL_TENSOR], dict[int, Tensor]] = {}
+
+    # caches one half of a merged weight until the other half arrives, then yields [a | b]
+    def _merge_shared_input(self, new_name: str, data: Tensor, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # a pre-quantized checkpoint carries per-tensor scales, which one merged weight cannot hold
+        if not self._SHARED_INPUT_MERGES or bid is None or self.hparams.get("quantization_config") or self._fp8_e4m3_preserved or self.hadamard_folded_names() or hasattr(data, "get_lora_A_B"):
+            yield (new_name, data)
+            return
+
+        layer_types = self.hparams.get("layer_types")
+        if layer_types is None:
+            is_linear = bid < self.hparams["num_hidden_layers"] and (bid + 1) % self.hparams.get("full_attention_interval", 4) != 0
+        else:
+            is_linear = bid < len(layer_types) and layer_types[bid] == "linear_attention"
+
+        match = next(((merged, parts.index(part)) for merged, parts, linear_only in self._SHARED_INPUT_MERGES
+                      if merged in gguf.MODEL_TENSORS[self.model_arch]
+                      for part in parts
+                      if (is_linear or not linear_only) and self.match_model_tensor_name(new_name, part, bid)), None)
+        if match is None:
+            yield (new_name, data)
+            return
+
+        merged, idx = match
+        halves = self._merge_parts.setdefault((bid, merged), {})
+        halves[idx] = data
+        if len(halves) == 2:
+            del self._merge_parts[(bid, merged)]
+            yield (self.format_tensor_name(merged, bid), torch.cat([halves[0], halves[1]], dim=0))
+
+    def prepare_tensors(self):
+        super().prepare_tensors()
+        if self._merge_parts:
+            raise ValueError(f"merged weights missing a half: {sorted((bid, t.name) for bid, t in self._merge_parts)}")
+
     @staticmethod
     def _reorder_v_heads(tensor: Tensor, dim: int, num_k_heads: int, num_v_per_k: int, head_dim: int) -> Tensor:
         """Reorder V heads from grouped (by K head) to tiled order along the given dimension."""
@@ -672,7 +711,8 @@ class _LinearAttentionVReorderBase(Qwen3NextModel):
                     # Out projection weight: reorder columns (input dimension)
                     data_torch = self._reorder_v_heads(data_torch, 1, num_k_heads, num_v_per_k, head_v_dim)
 
-        yield from super().modify_tensors(data_torch, name, bid)
+        for new_name, data in super().modify_tensors(data_torch, name, bid):
+            yield from self._merge_shared_input(new_name, data, bid)
 
 
 class _Qwen35MRopeMixin:
@@ -696,6 +736,11 @@ class _Qwen35MRopeMixin:
 class Qwen3_5TextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
     model_arch = gguf.MODEL_ARCH.QWEN35
 
+    _SHARED_INPUT_MERGES = (
+        (gguf.MODEL_TENSOR.ATTN_QKV_GATE,      (gguf.MODEL_TENSOR.ATTN_QKV, gguf.MODEL_TENSOR.ATTN_GATE), True),
+        (gguf.MODEL_TENSOR.SSM_BETA_ALPHA_CAT, (gguf.MODEL_TENSOR.SSM_BETA, gguf.MODEL_TENSOR.SSM_ALPHA), True),
+    )
+
     def __init__(self, dir_model, *args, **kwargs):
         # Inner TextModel does not own mtp.*. Set no_mtp before mixin bumps block_count.
         hparams = kwargs.pop("hparams", None)

conversion/qwen4exp.py

@@ -41,8 +41,6 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
         self._ple_shards: dict[int, str] = {}
         self._ple_row_dim: int | None = None
         self._mtp_fc: dict[str, Tensor] = {}
-        # halves of the merged weights, keyed by (bid, merged tensor), see _merge_shared_input
-        self._merge_parts: dict[tuple[int, gguf.MODEL_TENSOR], dict[int, Tensor]] = {}
 
     @classmethod
     def filter_tensors(cls, item):
@@ -171,8 +169,7 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
             yield (self.map_tensor_name(name), data_torch.squeeze())
             return
 
-        for new_name, data in super().modify_tensors(data_torch, name, bid):
-            yield from self._merge_shared_input(new_name, data, bid)
+        yield from super().modify_tensors(data_torch, name, bid)
 
     # matmuls that read the same input, written as one weight [a | b] so the runtime does one matmul:
     # (merged tensor, (a, b), linear-attention layers only)
@@ -183,30 +180,6 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
         (gguf.MODEL_TENSOR.SSM_BETA_ALPHA_CAT,  (gguf.MODEL_TENSOR.SSM_BETA,     gguf.MODEL_TENSOR.SSM_ALPHA),      True),
     )
 
-    # caches one half of a merged weight until the other half arrives, then yields [a | b]
-    def _merge_shared_input(self, new_name: str, data: Tensor, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # a pre-quantized checkpoint carries per-tensor scales, which one merged weight cannot hold
-        if bid is None or self.hparams.get("quantization_config"):
-            yield (new_name, data)
-            return
-
-        layer_types = self.hparams.get("layer_types") or []
-        is_linear = bid < len(layer_types) and layer_types[bid] == "linear_attention"
-
-        match = next(((merged, parts.index(part)) for merged, parts, linear_only in self._SHARED_INPUT_MERGES
-                      for part in parts
-                      if (is_linear or not linear_only) and self.match_model_tensor_name(new_name, part, bid)), None)
-        if match is None:
-            yield (new_name, data)
-            return
-
-        merged, idx = match
-        halves = self._merge_parts.setdefault((bid, merged), {})
-        halves[idx] = data
-        if len(halves) == 2:
-            del self._merge_parts[(bid, merged)]
-            yield (self.format_tensor_name(merged, bid), torch.cat([halves[0], halves[1]], dim=0))
-
     # the shards concatenate into a tensor of well over 100 GB
     # use LazyChunkedTensor here, a single shard resident at a time
     def _place_ple_shard(self, data_torch: Tensor, name: str) -> Iterable[tuple[str, Tensor]]:
@@ -251,8 +224,6 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
         super().prepare_tensors()
         if self._mtp_fc:
             raise ValueError(f"MTP projection missing its other half: {sorted(self._mtp_fc)}")
-        if self._merge_parts:
-            raise ValueError(f"merged weights missing a half: {sorted((bid, t.name) for bid, t in self._merge_parts)}")
         n_parts = self.hparams.get("split_ngram_parts", 0)
         if self._ple_shards and len(self._ple_shards) != n_parts:
             raise ValueError(

docs/backend/snapdragon/developer.md

@@ -60,7 +60,7 @@ Implementation: [`htp-repack.cpp`](../../../ggml/src/ggml-hexagon/htp-repack.cpp
 ### Tile Dimensions and Grid Layout
 
 Each tile represents a submatrix of `N_rows = 32` rows and `K_cols` columns:
-- Standard types (`Q4_0`, `Q4_1`, `Q8_0`, `IQ4_NL`, `MXFP4`, and K-quants): `K_cols = 32` (`HTP_TENSOR_TILE_N_COLS = 32`).
+- Standard types (`Q4_0`, `Q4_1`, `Q5_0`, `Q5_1`, `Q8_0`, `IQ4_NL`, `MXFP4`, and K-quants): `K_cols = 32` (`HTP_TENSOR_TILE_N_COLS = 32`).
 - `Q2_0`: `K_cols = 64` (`HTP_TENSOR_TILE_K_Q2_0 = 64`).
 - `Q1_0`: `K_cols = 128` (`HTP_TENSOR_TILE_K_Q1_0 = 128`).
 
@@ -123,6 +123,17 @@ Block-32 formats use a tile size of 32 rows x 32 columns (`K_cols = 32`):
   - **Scales** (bytes 512..639, 128 bytes):
     32 pairs of FP16 values (`ggml_half`): `scale[2 * row + 0] = d` (scale) and `scale[2 * row + 1] = m` (min) for each row `row = 0..31`.
 
+- **Q5_0** (704 bytes per tile, `HTP_TENSOR_TILE_SIZE_Q5_0`):
+  - **Quants** (bytes 0..511): same 4-bit nibble transposition and storage layout as `Q4_0`, holding the low 4 bits.
+  - **5th-bit plane** (bytes 512..639): byte `32*c + row` (`c = 0..3`) holds a flag for `k = 2*(4i + c)` at bit `2i` and for `k = 2*(4i + c) + 1` at bit `2i + 1`, for `i = 0..3`. The flag is set when the 5th bit of the quant is clear.
+  - **Scales** (bytes 640..703, 64 bytes): 32 FP16 values, storing scale `d` for each row `row = 0..31`.
+  - **Dequantization**: `0xF0` is OR-ed into the nibble where the flag is set, giving the signed value `q - 16` (-16..15), scaled by `d`.
+
+- **Q5_1** (768 bytes per tile, `HTP_TENSOR_TILE_SIZE_Q5_1`, same layout as the Q5_K tile):
+  - **Quants** (bytes 0..511) and **scales** (bytes 512..639): same as `Q4_1`, holding the low 4 bits and `(d, m)` pairs.
+  - **5th-bit plane** (bytes 640..767): same as `Q5_K`; a bit is set when the 5th bit of the quant is set.
+  - Uses the `Q5_K` HVX/HMX kernels (Q8_1 activations with block sums).
+
 ### Q1_0 Format (1-Bit Quantization)
 
 `Q1_0` packs 128 columns per tile (`K_cols = 128`, `HTP_TENSOR_TILE_K_Q1_0 = 128`):

docs/build.md

@@ -608,7 +608,7 @@ You can test with:
 If the following info is output on screen, you are using `llama.cpp` with the CANN backend:
 ```bash
 llm_load_tensors:       CANN model buffer size = 13313.00 MiB
-llama_new_context_with_model:       CANN compute buffer size =  1260.81 MiB
+llama_init_from_model:       CANN compute buffer size =  1260.81 MiB
 ```
 
 For detailed info, such as model/device supports, CANN install, please refer to [llama.cpp for CANN](./backend/CANN.md).

docs/ops.md

@@ -40,7 +40,6 @@ Legend:
 |          CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                           CUMSUM | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             DIAG | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                    DIAG_MASK_INF | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              DIV | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                     DSV4_HC_COMB | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                     DSV4_HC_POST | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |

docs/ops/BLAS.csv

@@ -8741,9 +8741,6 @@
 "BLAS","ROUND","type=f32,ne=[1024,1024,1,1]","support","0","no","BLAS"
 "BLAS","TRUNC","type=f32,ne=[7,1,5,3]","support","0","no","BLAS"
 "BLAS","TRUNC","type=f32,ne=[1024,1024,1,1]","support","0","no","BLAS"
-"BLAS","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","0","no","BLAS"
-"BLAS","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","0","no","BLAS"
-"BLAS","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","0","no","BLAS"
 "BLAS","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","BLAS"
 "BLAS","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","BLAS"
 "BLAS","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","BLAS"

docs/ops/CANN.csv

@@ -8764,9 +8764,6 @@
 "CANN0","ROUND","type=f32,ne=[1024,1024,1,1]","support","0","no","CANN"
 "CANN0","TRUNC","type=f32,ne=[7,1,5,3]","support","0","no","CANN"
 "CANN0","TRUNC","type=f32,ne=[1024,1024,1,1]","support","0","no","CANN"
-"CANN0","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","1","yes","CANN"
-"CANN0","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","1","yes","CANN"
-"CANN0","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","1","yes","CANN"
 "CANN0","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CANN"
 "CANN0","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CANN"
 "CANN0","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CANN"

docs/ops/CPU.csv

@@ -12962,9 +12962,6 @@
 "CPU","ROUND","type=f32,ne=[1024,1024,1,1]","support","1","yes","CPU"
 "CPU","TRUNC","type=f32,ne=[7,1,5,3]","support","1","yes","CPU"
 "CPU","TRUNC","type=f32,ne=[1024,1024,1,1]","support","1","yes","CPU"
-"CPU","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","1","yes","CPU"
-"CPU","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","1","yes","CPU"
-"CPU","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","1","yes","CPU"
 "CPU","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CPU"
 "CPU","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CPU"
 "CPU","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CPU"

docs/ops/CUDA.csv

@@ -8725,9 +8725,6 @@
 "CUDA0","ROUND","type=f32,ne=[1024,1024,1,1]","support","1","yes","CUDA"
 "CUDA0","TRUNC","type=f32,ne=[7,1,5,3]","support","1","yes","CUDA"
 "CUDA0","TRUNC","type=f32,ne=[1024,1024,1,1]","support","1","yes","CUDA"
-"CUDA0","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","1","yes","CUDA"
-"CUDA0","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","1","yes","CUDA"
-"CUDA0","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","1","yes","CUDA"
 "CUDA0","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CUDA"
 "CUDA0","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CUDA"
 "CUDA0","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","CUDA"

docs/ops/ET.csv

@@ -9637,9 +9637,6 @@
 "ET","ROUND","type=f32,ne=[1024,1024,1,1]","support","1","yes","ET"
 "ET","TRUNC","type=f32,ne=[7,1,5,3]","support","0","no","ET"
 "ET","TRUNC","type=f32,ne=[1024,1024,1,1]","support","1","yes","ET"
-"ET","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","0","no","ET"
-"ET","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","0","no","ET"
-"ET","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","0","no","ET"
 "ET","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","ET"
 "ET","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","ET"
 "ET","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","ET"

docs/ops/Hexagon.csv

@@ -14564,9 +14564,6 @@
 "HTP0","ROUND","type=f32,ne=[1024,1024,1,1]","support","0","no","HTP"
 "HTP0","TRUNC","type=f32,ne=[7,1,5,3]","support","0","no","HTP"
 "HTP0","TRUNC","type=f32,ne=[1024,1024,1,1]","support","0","no","HTP"
-"HTP0","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","0","no","HTP"
-"HTP0","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","0","no","HTP"
-"HTP0","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","0","no","HTP"
 "HTP0","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","HTP"
 "HTP0","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","HTP"
 "HTP0","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","HTP"

docs/ops/Metal.csv

@@ -9628,9 +9628,6 @@
 "MTL0","ROUND","type=f32,ne=[1024,1024,1,1]","support","1","yes","MTL"
 "MTL0","TRUNC","type=f32,ne=[7,1,5,3]","support","1","yes","MTL"
 "MTL0","TRUNC","type=f32,ne=[1024,1024,1,1]","support","1","yes","MTL"
-"MTL0","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","0","no","MTL"
-"MTL0","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","0","no","MTL"
-"MTL0","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","0","no","MTL"
 "MTL0","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","MTL"
 "MTL0","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","MTL"
 "MTL0","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","MTL"

docs/ops/OpenCL.csv

@@ -14887,9 +14887,6 @@
 "OpenCL","ROUND","type=f32,ne=[1024,1024,1,1]","support","1","yes","OpenCL"
 "OpenCL","TRUNC","type=f32,ne=[7,1,5,3]","support","1","yes","OpenCL"
 "OpenCL","TRUNC","type=f32,ne=[1024,1024,1,1]","support","1","yes","OpenCL"
-"OpenCL","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","1","yes","OpenCL"
-"OpenCL","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","1","yes","OpenCL"
-"OpenCL","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","0","no","OpenCL"
 "OpenCL","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","OpenCL"
 "OpenCL","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","OpenCL"
 "OpenCL","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","OpenCL"

docs/ops/SYCL.csv

@@ -14534,9 +14534,6 @@
 "SYCL0","ROUND","type=f32,ne=[1024,1024,1,1]","support","1","yes","SYCL"
 "SYCL0","TRUNC","type=f32,ne=[7,1,5,3]","support","1","yes","SYCL"
 "SYCL0","TRUNC","type=f32,ne=[1024,1024,1,1]","support","1","yes","SYCL"
-"SYCL0","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","1","yes","SYCL"
-"SYCL0","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","1","yes","SYCL"
-"SYCL0","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","1","yes","SYCL"
 "SYCL0","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","SYCL"
 "SYCL0","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","SYCL"
 "SYCL0","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","SYCL"

docs/ops/Vulkan.csv

@@ -12313,9 +12313,6 @@
 "Vulkan0","ROUND","type=f32,ne=[1024,1024,1,1]","support","1","yes","Vulkan"
 "Vulkan0","TRUNC","type=f32,ne=[7,1,5,3]","support","1","yes","Vulkan"
 "Vulkan0","TRUNC","type=f32,ne=[1024,1024,1,1]","support","1","yes","Vulkan"
-"Vulkan0","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","1","yes","Vulkan"
-"Vulkan0","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","1","yes","Vulkan"
-"Vulkan0","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","1","yes","Vulkan"
 "Vulkan0","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","Vulkan"
 "Vulkan0","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","Vulkan"
 "Vulkan0","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","Vulkan"

docs/ops/WebGPU.csv

@@ -10560,9 +10560,6 @@
 "WebGPU: WebGPU","ROUND","type=f32,ne=[1024,1024,1,1]","support","1","yes","WebGPU"
 "WebGPU: WebGPU","TRUNC","type=f32,ne=[7,1,5,3]","support","1","yes","WebGPU"
 "WebGPU: WebGPU","TRUNC","type=f32,ne=[1024,1024,1,1]","support","1","yes","WebGPU"
-"WebGPU: WebGPU","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","0","no","WebGPU"
-"WebGPU: WebGPU","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","0","no","WebGPU"
-"WebGPU: WebGPU","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","0","no","WebGPU"
 "WebGPU: WebGPU","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","WebGPU"
 "WebGPU: WebGPU","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","WebGPU"
 "WebGPU: WebGPU","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","1","yes","WebGPU"

docs/ops/ZenDNN.csv

@@ -9319,9 +9319,6 @@
 "ZenDNN","ROUND","type=f32,ne=[1024,1024,1,1]","support","0","no","ZenDNN"
 "ZenDNN","TRUNC","type=f32,ne=[7,1,5,3]","support","0","no","ZenDNN"
 "ZenDNN","TRUNC","type=f32,ne=[1024,1024,1,1]","support","0","no","ZenDNN"
-"ZenDNN","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","0","no","ZenDNN"
-"ZenDNN","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","0","no","ZenDNN"
-"ZenDNN","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","0","no","ZenDNN"
 "ZenDNN","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","ZenDNN"
 "ZenDNN","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","ZenDNN"
 "ZenDNN","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","ZenDNN"

docs/ops/zDNN.csv

@@ -14550,9 +14550,6 @@
 "zDNN","ROUND","type=f32,ne=[1024,1024,1,1]","support","0","no","zDNN"
 "zDNN","TRUNC","type=f32,ne=[7,1,5,3]","support","0","no","zDNN"
 "zDNN","TRUNC","type=f32,ne=[1024,1024,1,1]","support","0","no","zDNN"
-"zDNN","DIAG_MASK_INF","type=f32,ne=[10,10,1,1],n_past=5","support","0","no","zDNN"
-"zDNN","DIAG_MASK_INF","type=f32,ne=[10,10,3,1],n_past=5","support","0","no","zDNN"
-"zDNN","DIAG_MASK_INF","type=f32,ne=[10,10,3,2],n_past=5","support","0","no","zDNN"
 "zDNN","SOFT_MAX","type=f32,ne=[16,16,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","zDNN"
 "zDNN","SOFT_MAX","type=f32,ne=[15,15,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","zDNN"
 "zDNN","SOFT_MAX","type=f32,ne=[16,1024,1,1],mask=0,sinks=0,m_prec=f32,nr23=[1,1],scale=1.000000,max_bias=0.000000,inplace=0","support","0","no","zDNN"

examples/eval-callback/README.md

@@ -20,12 +20,12 @@ Will print:
 ```shell
 llm_load_tensors: offloaded 33/33 layers to GPU
 ...
-llama_new_context_with_model: n_ctx      = 512
+llama_init_from_model: n_ctx      = 512
 ...
-llama_new_context_with_model:      CUDA0 compute buffer size =   105.00 MiB
-llama_new_context_with_model:  CUDA_Host compute buffer size =     6.01 MiB
-llama_new_context_with_model: graph nodes  = 1225
-llama_new_context_with_model: graph splits = 2
+llama_init_from_model:      CUDA0 compute buffer size =   105.00 MiB
+llama_init_from_model:  CUDA_Host compute buffer size =     6.01 MiB
+llama_init_from_model: graph nodes  = 1225
+llama_init_from_model: graph splits = 2
 ggml_debug:                 inp_embd = (f32)   GET_ROWS(token_embd.weight{2560, 51200, 1, 1}, inp_tokens{1, 1, 1, 1}}) = {2560, 1, 1, 1}
                                      [
                                       [

examples/llama.android/lib/src/main/cpp/ai_chat.cpp

@@ -99,7 +99,7 @@ static llama_context *init_context(llama_model *model, const int n_ctx = DEFAULT
     ctx_params.n_threads_batch = n_threads;
     auto *context = llama_init_from_model(g_model, ctx_params);
     if (context == nullptr) {
-        LOGe("%s: llama_new_context_with_model() returned null)", __func__);
+        LOGe("%s: llama_init_from_model() returned null)", __func__);
     }
     return context;
 }

ggml/include/ggml-rpc.h

@@ -6,26 +6,31 @@
 extern "C" {
 #endif
 
-#define RPC_PROTO_MAJOR_VERSION    9
+#define RPC_PROTO_MAJOR_VERSION    10
 #define RPC_PROTO_MINOR_VERSION    0
 #define RPC_PROTO_PATCH_VERSION    0
 
 #ifdef  __cplusplus
-static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
 #endif
 
 #define GGML_RPC_MAX_SERVERS       16
 
 // backend API
 GGML_BACKEND_API ggml_backend_t ggml_backend_rpc_init(const char * endpoint, uint32_t device);
+
 GGML_BACKEND_API bool ggml_backend_is_rpc(ggml_backend_t backend);
 
 GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint, uint32_t device);
 
 GGML_BACKEND_API void ggml_backend_rpc_get_device_memory(const char * endpoint, uint32_t device, size_t * free, size_t * total);
 
-GGML_BACKEND_API void ggml_backend_rpc_start_server(const char * endpoint, const char * cache_dir,
-                                                    size_t n_threads, size_t n_devices, ggml_backend_dev_t * devices);
+GGML_BACKEND_API void ggml_backend_rpc_start_server(
+        const char * endpoint,
+        const char * cache_dir,
+        size_t n_threads,
+        size_t n_devices,
+        ggml_backend_dev_t * devices);
 
 GGML_BACKEND_API ggml_backend_reg_t ggml_backend_rpc_reg(void);
 GGML_BACKEND_API ggml_backend_reg_t ggml_backend_rpc_add_server(const char * endpoint);

ggml/include/ggml.h

@@ -1,178 +1,5 @@
 #pragma once
 
-//
-// GGML Tensor Library
-//
-// This documentation is still a work in progress.
-// If you wish some specific topics to be covered, feel free to drop a comment:
-//
-//   https://github.com/ggml-org/whisper.cpp/issues/40
-//
-// ## Overview
-//
-// This library implements:
-//
-//  - a set of tensor operations
-//  - automatic differentiation
-//  - basic optimization algorithms
-//
-// The aim of this library is to provide a minimalistic approach for various machine learning tasks. This includes,
-// but is not limited to, the following:
-//
-//  - linear regression
-//  - support vector machines
-//  - neural networks
-//
-// The library allows the user to define a certain function using the available tensor operations. This function
-// definition is represented internally via a computation graph. Each tensor operation in the function definition
-// corresponds to a node in the graph. Having the computation graph defined, the user can choose to compute the
-// function's value and/or its gradient with respect to the input variables. Optionally, the function can be optimized
-// using one of the available optimization algorithms.
-//
-// For example, here we define the function: f(x) = a*x^2 + b
-//
-//   {
-//       struct ggml_init_params params = {
-//           .mem_size   = 16*1024*1024,
-//           .mem_buffer = NULL,
-//       };
-//
-//       // memory allocation happens here
-//       struct ggml_context * ctx = ggml_init(params);
-//
-//       struct ggml_tensor * x = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
-//
-//       ggml_set_param(ctx, x); // x is an input variable
-//
-//       struct ggml_tensor * a  = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
-//       struct ggml_tensor * b  = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
-//       struct ggml_tensor * x2 = ggml_mul(ctx, x, x);
-//       struct ggml_tensor * f  = ggml_add(ctx, ggml_mul(ctx, a, x2), b);
-//
-//       ...
-//   }
-//
-// Notice that the function definition above does not involve any actual computation. The computation is performed only
-// when the user explicitly requests it. For example, to compute the function's value at x = 2.0:
-//
-//   {
-//       ...
-//
-//       struct ggml_cgraph * gf = ggml_new_graph(ctx);
-//       ggml_build_forward_expand(gf, f);
-//
-//       // set the input variable and parameter values
-//       ggml_set_f32(x, 2.0f);
-//       ggml_set_f32(a, 3.0f);
-//       ggml_set_f32(b, 4.0f);
-//
-//       ggml_graph_compute_with_ctx(ctx, &gf, n_threads);
-//
-//       printf("f = %f\n", ggml_get_f32_1d(f, 0));
-//
-//       ...
-//   }
-//
-// The actual computation is performed in the ggml_graph_compute() function.
-//
-// The ggml_new_tensor_...() functions create new tensors. They are allocated in the memory buffer provided to the
-// ggml_init() function. You have to be careful not to exceed the memory buffer size. Therefore, you have to know
-// in advance how much memory you need for your computation. Alternatively, you can allocate a large enough memory
-// and after defining the computation graph, call the ggml_used_mem() function to find out how much memory was
-// actually needed.
-//
-// The ggml_set_param() function marks a tensor as an input variable. This is used by the automatic
-// differentiation and optimization algorithms.
-//
-// The described approach allows to define the function graph once and then compute its forward or backward graphs
-// multiple times. All computations will use the same memory buffer allocated in the ggml_init() function. This way
-// the user can avoid the memory allocation overhead at runtime.
-//
-// The library supports multi-dimensional tensors - up to 4 dimensions. The FP16 and FP32 data types are first class
-// citizens, but in theory the library can be extended to support FP8 and integer data types.
-//
-// Each tensor operation produces a new tensor. Initially the library was envisioned to support only the use of unary
-// and binary operations. Most of the available operations fall into one of these two categories. With time, it became
-// clear that the library needs to support more complex operations. The way to support these operations is not clear
-// yet, but a few examples are demonstrated in the following operations:
-//
-//   - ggml_permute()
-//   - ggml_conv_1d_1s()
-//   - ggml_conv_1d_2s()
-//
-// For each tensor operator, the library implements a forward and backward computation function. The forward function
-// computes the output tensor value given the input tensor values. The backward function computes the adjoint of the
-// input tensors given the adjoint of the output tensor. For a detailed explanation of what this means, take a
-// calculus class, or watch the following video:
-//
-//   What is Automatic Differentiation?
-//   https://www.youtube.com/watch?v=wG_nF1awSSY
-//
-//
-// ## Tensor data (struct ggml_tensor)
-//
-// The tensors are stored in memory via the ggml_tensor struct. The structure provides information about the size of
-// the tensor, the data type, and the memory buffer where the tensor data is stored. Additionally, it contains
-// pointers to the "source" tensors - i.e. the tensors that were used to compute the current tensor. For example:
-//
-//   {
-//       struct ggml_tensor * c = ggml_add(ctx, a, b);
-//
-//       assert(c->src[0] == a);
-//       assert(c->src[1] == b);
-//   }
-//
-// The multi-dimensional tensors are stored in row-major order. The ggml_tensor struct contains fields for the
-// number of elements in each dimension ("ne") as well as the number of bytes ("nb", a.k.a. stride). This allows
-// to store tensors that are not contiguous in memory, which is useful for operations such as transposition and
-// permutation. All tensor operations have to take the stride into account and not assume that the tensor is
-// contiguous in memory.
-//
-// The data of the tensor is accessed via the "data" pointer. For example:
-//
-//   {
-//       const int nx = 2;
-//       const int ny = 3;
-//
-//       struct ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, nx, ny);
-//
-//       for (int y = 0; y < ny; y++) {
-//           for (int x = 0; x < nx; x++) {
-//               *(float *) ((char *) a->data + y*a->nb[1] + x*a->nb[0]) = x + y;
-//           }
-//       }
-//
-//       ...
-//   }
-//
-// Alternatively, there are helper functions, such as ggml_get_f32_1d() and ggml_set_f32_1d() that can be used.
-//
-// ## The matrix multiplication operator (ggml_mul_mat)
-//
-// TODO
-//
-//
-// ## Multi-threading
-//
-// TODO
-//
-//
-// ## Overview of ggml.c
-//
-// TODO
-//
-//
-// ## SIMD optimizations
-//
-// TODO
-//
-//
-// ## Debugging ggml
-//
-// TODO
-//
-//
-
 #ifdef GGML_SHARED
 #    if defined(_WIN32) && !defined(__MINGW32__)
 #        ifdef GGML_BUILD
@@ -445,7 +272,6 @@ extern "C" {
     // the precision parameters are stored as ggml_tensor.op_params to the respective ops
     enum ggml_prec {
         GGML_PREC_UNDEFINED = 0,
-        GGML_PREC_DEFAULT   = 0,  // note: deprecated, use GGML_PREC_UNDEFINED
         GGML_PREC_F32       = 10,
         GGML_PREC_BF16      = 15,
         GGML_PREC_F16       = 20,
@@ -499,7 +325,6 @@ extern "C" {
         GGML_OP_DUP,
         GGML_OP_ADD,
         GGML_OP_ADD_ID,
-        GGML_OP_ADD1,
         GGML_OP_ACC,
         GGML_OP_SUB,
         GGML_OP_MUL,
@@ -541,8 +366,6 @@ extern "C" {
         GGML_OP_GET_ROWS_BACK,
         GGML_OP_SET_ROWS,
         GGML_OP_DIAG,
-        GGML_OP_DIAG_MASK_INF,
-        GGML_OP_DIAG_MASK_ZERO,
         GGML_OP_SOFT_MAX,
         GGML_OP_SOFT_MAX_BACK,
         GGML_OP_ROPE,
@@ -765,10 +588,6 @@ extern "C" {
     GGML_API size_t  ggml_type_size(enum ggml_type type);             // size in bytes for all elements in a block
     GGML_API size_t  ggml_row_size (enum ggml_type type, int64_t ne); // size in bytes for all elements in a row
 
-    GGML_DEPRECATED(
-    GGML_API double ggml_type_sizef(enum ggml_type type), // ggml_type_size()/ggml_blck_size() as float
-    "use ggml_row_size() instead");
-
     GGML_API const char * ggml_type_name(enum ggml_type type);
     GGML_API const char * ggml_op_name  (enum ggml_op   op);
     GGML_API const char * ggml_op_symbol(enum ggml_op   op);
@@ -937,18 +756,6 @@ extern "C" {
             struct ggml_tensor  * b,
             struct ggml_tensor  * ids);
 
-    GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_add1(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            struct ggml_tensor  * b),
-        "use ggml_add instead");
-
-    GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_add1_inplace(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            struct ggml_tensor  * b),
-        "use ggml_add_inplace instead");
-
     // dst = a
     // view(dst, nb1, nb2, nb3, offset) += b
     // return dst
@@ -1500,13 +1307,6 @@ extern "C" {
             struct ggml_tensor  * a,
             struct ggml_tensor  * b);
 
-    // change the precision of a matrix multiplication
-    // set to GGML_PREC_F32 for higher precision (useful for phi-2)
-    GGML_DEPRECATED(GGML_API void ggml_mul_mat_set_prec(
-            struct ggml_tensor * a,
-            enum ggml_prec       prec),
-        "use ggml_prec_set_acc() instead");
-
     // change the hint of a matrix multiplication
     GGML_API void ggml_mul_mat_set_hint(
             struct ggml_tensor * a,
@@ -1778,30 +1578,6 @@ extern "C" {
         struct ggml_context     * ctx,
         struct ggml_tensor      * a);
 
-    // set elements above the diagonal to -INF
-    GGML_API struct ggml_tensor * ggml_diag_mask_inf(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            int                   n_past);
-
-    // in-place, returns view(a)
-    GGML_API struct ggml_tensor * ggml_diag_mask_inf_inplace(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            int                   n_past);
-
-    // set elements above the diagonal to 0
-    GGML_API struct ggml_tensor * ggml_diag_mask_zero(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            int                   n_past);
-
-    // in-place, returns view(a)
-    GGML_API struct ggml_tensor * ggml_diag_mask_zero_inplace(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            int                   n_past);
-
     GGML_API struct ggml_tensor * ggml_clamp(
             struct ggml_context * ctx,
             struct ggml_tensor  * a,
@@ -2005,36 +1781,6 @@ extern "C" {
             float                 beta_fast,
             float                 beta_slow);
 
-    GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_rope_custom(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            struct ggml_tensor  * b,
-            int                   n_dims,
-            int                   mode,
-            int                   n_ctx_orig,
-            float                 freq_base,
-            float                 freq_scale,
-            float                 ext_factor,
-            float                 attn_factor,
-            float                 beta_fast,
-            float                 beta_slow),
-        "use ggml_rope_ext instead");
-
-    GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_rope_custom_inplace(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            struct ggml_tensor  * b,
-            int                   n_dims,
-            int                   mode,
-            int                   n_ctx_orig,
-            float                 freq_base,
-            float                 freq_scale,
-            float                 ext_factor,
-            float                 attn_factor,
-            float                 beta_fast,
-            float                 beta_slow),
-        "use ggml_rope_ext_inplace instead");
-
     // rotary position embedding backward, i.e compute dx from dy
     // a - dy
     GGML_API struct ggml_tensor * ggml_rope_ext_back(
@@ -2340,7 +2086,7 @@ extern "C" {
         GGML_SCALE_MODE_BILINEAR = 1,
         GGML_SCALE_MODE_BICUBIC  = 2,
 
-        GGML_SCALE_MODE_COUNT
+        GGML_SCALE_MODE_COUNT = 3
     };
 
     enum ggml_scale_flag {
@@ -2356,18 +2102,6 @@ extern "C" {
             int                   scale_factor,
             enum ggml_scale_mode  mode);
 
-    // interpolate
-    // interpolate scale to specified dimensions
-    GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_upscale_ext(
-            struct ggml_context * ctx,
-            struct ggml_tensor  * a,
-            int                   ne0,
-            int                   ne1,
-            int                   ne2,
-            int                   ne3,
-            enum ggml_scale_mode  mode),
-        "use ggml_interpolate instead");
-
     // Up- or downsamples the input to the specified size.
     // 2D scale modes (eg. bilinear) are applied to the first two dimensions.
     GGML_API struct ggml_tensor * ggml_interpolate(
@@ -2536,11 +2270,6 @@ extern "C" {
             float                 max_bias,
             float                 logit_softcap);
 
-    GGML_DEPRECATED(GGML_API void ggml_flash_attn_ext_set_prec(
-            struct ggml_tensor * a,
-            enum ggml_prec       prec),
-        "use ggml_prec_set_acc() instead");
-
     GGML_API enum ggml_prec ggml_flash_attn_ext_get_prec(
             const struct ggml_tensor * a);
 

ggml/src/ggml-alloc.c

@@ -23,11 +23,8 @@ bool ggml_op_can_inplace(enum ggml_op op) {
     switch (op) {
         case GGML_OP_FILL:
         case GGML_OP_SCALE:
-        case GGML_OP_DIAG_MASK_ZERO:
-        case GGML_OP_DIAG_MASK_INF:
         case GGML_OP_ADD:
         case GGML_OP_ADD_ID:
-        case GGML_OP_ADD1:
         case GGML_OP_SUB:
         case GGML_OP_MUL:
         case GGML_OP_DIV:

ggml/src/ggml-backend-meta.cpp

@@ -939,7 +939,6 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             case GGML_OP_ADD_ID: {
                 split_state = handle_bin_bcast(src_ss);
             } break;
-            case GGML_OP_ADD1:
             case GGML_OP_ACC: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ true);
             } break;
@@ -1020,9 +1019,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             case GGML_OP_SET_ROWS: {
                 split_state = handle_set_rows(src_ss);
             } break;
-            case GGML_OP_DIAG:
-            case GGML_OP_DIAG_MASK_INF:
-            case GGML_OP_DIAG_MASK_ZERO: {
+            case GGML_OP_DIAG: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ true);
             } break;
             case GGML_OP_SOFT_MAX:

ggml/src/ggml-backend.cpp

@@ -2172,8 +2172,13 @@ static void ggml_backend_sched_prefetch_inputs(ggml_backend_sched_t sched, int s
             continue;
         }
 
-        // not if this split computes or writes the input
+        // only results of earlier splits, the copy callback of this split runs later and can write other tensors (e.g. the MoE cache)
         const struct ggml_tensor * base = input->view_src ? input->view_src : input;
+        if (base->buffer == NULL || ggml_backend_buffer_get_usage(base->buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
+            continue;
+        }
+
+        // not if this split computes or writes the input
         bool written = false;
         for (int j = 0; j < split->graph.n_nodes && !written; j++) {
             const struct ggml_tensor * node = split->graph.nodes[j];

ggml/src/ggml-cann/aclnn_ops.cpp

@@ -1347,31 +1347,6 @@ void ggml_cann_rms_norm(ggml_backend_cann_context & ctx, ggml_tensor * dst) {
     GGML_CANN_CALL_ACLNN_OP(ctx, RmsNorm, acl_src.get(), acl_gamma.get(), eps, acl_dst.get(), acl_rstd.get());
 }
 
-// TODO: performace is low.
-void ggml_cann_diag_mask(ggml_backend_cann_context & ctx, ggml_tensor * dst, float value) {
-    ggml_tensor * src = dst->src[0];
-
-    acl_tensor_ptr acl_src = ggml_cann_create_tensor(src);
-    acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst);
-
-    const int n_past = ((int32_t *) dst->op_params)[0];
-
-    ggml_cann_pool_alloc one_tensor_allocator(ctx.pool(), ggml_nbytes(src));
-    void *               buffer = one_tensor_allocator.get();
-
-    acl_tensor_ptr mask_tensor = ggml_cann_create_tensor(buffer, ggml_cann_type_mapping(src->type),
-                                                         ggml_type_size(src->type), src->ne, src->nb, GGML_MAX_DIMS);
-
-    aclnn_fill_scalar(ctx, value, mask_tensor.get());
-
-    float          alphaValue = 1.0f;
-    acl_scalar_ptr alpha      = ggml_cann_create_scalar(&alphaValue, aclDataType::ACL_FLOAT);
-
-    GGML_CANN_CALL_ACLNN_OP(ctx, InplaceTriu, mask_tensor.get(), n_past + 1);
-    GGML_CANN_CALL_ACLNN_OP(ctx, Tril, acl_src.get(), n_past + 1, acl_dst.get());
-    GGML_CANN_CALL_ACLNN_OP(ctx, InplaceAdd, acl_dst.get(), mask_tensor.get(), alpha.get());
-}
-
 /**
  * @brief Permutes the dimensions of a tensor according to a specified order.
  *

ggml/src/ggml-cann/aclnn_ops.h

@@ -466,21 +466,6 @@ void ggml_cann_dup(ggml_backend_cann_context & ctx, ggml_tensor * dst);
  */
 void ggml_cann_rms_norm(ggml_backend_cann_context & ctx, ggml_tensor * dst);
 
-/**
- * @brief   Applies a diagonal mask to the tensor with a specified value.
- *
- * @details This function creates a mask tensor filled with ones, then applies
- *          an upper triangular and lower triangular operation to it based on
- *          the number of past elements specified. Afterward, it adds the masked
- *          tensor to the destination tensor in-place.
- *
- * @param ctx The backend CANN context used for operations.
- * @param dst The destination tensor where the result will be stored. dst->op is
- *            `GGML_OP_DIAG_MASK`
- * @param value The value to use for masking.
- */
-void ggml_cann_diag_mask(ggml_backend_cann_context & ctx, ggml_tensor * dst, float value);
-
 /**
  * @brief   Performs an image-to-column transformation on the input tensor.
  *

ggml/src/ggml-cann/ggml-cann.cpp

@@ -1793,7 +1793,6 @@ static bool ggml_cann_compute_forward(ggml_backend_cann_context & ctx, struct gg
             ggml_cann_dup(ctx, dst);
             break;
         case GGML_OP_ADD:
-        case GGML_OP_ADD1:
             ggml_cann_binary_op<aclnn_add>(ctx, dst);
             break;
         case GGML_OP_SUB:
@@ -1958,9 +1957,6 @@ static bool ggml_cann_compute_forward(ggml_backend_cann_context & ctx, struct gg
         case GGML_OP_PERMUTE:
         case GGML_OP_TRANSPOSE:
             break;
-        case GGML_OP_DIAG_MASK_INF:
-            ggml_cann_diag_mask(ctx, dst, -INFINITY);
-            break;
         case GGML_OP_SOFT_MAX:
             ggml_cann_softmax(ctx, dst);
             break;
@@ -2618,7 +2614,6 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
         case GGML_OP_TRANSPOSE:
         case GGML_OP_NORM:
         case GGML_OP_ADD:
-        case GGML_OP_ADD1:
         case GGML_OP_SUB:
         case GGML_OP_MUL:
         case GGML_OP_DIV:
@@ -2626,7 +2621,6 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
         case GGML_OP_SQR:
         case GGML_OP_SQRT:
         case GGML_OP_CLAMP:
-        case GGML_OP_DIAG_MASK_INF:
         case GGML_OP_SUM_ROWS:
         case GGML_OP_ARGSORT:
         case GGML_OP_ACC:

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

@@ -1782,10 +1782,6 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_add_id(params, tensor);
             } break;
-        case GGML_OP_ADD1:
-            {
-                ggml_compute_forward_add1(params, tensor);
-            } break;
         case GGML_OP_ACC:
             {
                 ggml_compute_forward_acc(params, tensor);
@@ -1926,14 +1922,6 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_diag(params, tensor);
             } break;
-        case GGML_OP_DIAG_MASK_INF:
-            {
-                ggml_compute_forward_diag_mask_inf(params, tensor);
-            } break;
-        case GGML_OP_DIAG_MASK_ZERO:
-            {
-                ggml_compute_forward_diag_mask_zero(params, tensor);
-            } break;
         case GGML_OP_SOFT_MAX:
             {
                 ggml_compute_forward_soft_max(params, tensor);
@@ -2285,7 +2273,6 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_CONT:
         case GGML_OP_ADD:
         case GGML_OP_ADD_ID:
-        case GGML_OP_ADD1:
         case GGML_OP_ACC:
         case GGML_OP_CUMSUM:
         case GGML_OP_TRI:
@@ -2407,8 +2394,6 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
             {
                 n_tasks = 1;
             } break;
-        case GGML_OP_DIAG_MASK_ZERO:
-        case GGML_OP_DIAG_MASK_INF:
         case GGML_OP_SOFT_MAX_BACK:
         case GGML_OP_ROPE:
         case GGML_OP_ROPE_BACK:
@@ -2889,7 +2874,6 @@ struct ggml_cplan ggml_graph_plan(
                     } break;
                 case GGML_OP_ADD:
                 case GGML_OP_ADD_ID:
-                case GGML_OP_ADD1:
                     {
                         if (ggml_is_quantized(node->src[0]->type) || node->src[0]->type == GGML_TYPE_F8_E4M3) {
                             cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks;

ggml/src/ggml-cpu/ops.cpp

@@ -5704,94 +5704,6 @@ void ggml_compute_forward_diag(
     }
 }
 
-// ggml_compute_forward_diag_mask_inf
-
-static void ggml_compute_forward_diag_mask_f32(
-        const ggml_compute_params * params,
-        ggml_tensor * dst,
-        const float value) {
-
-    const ggml_tensor * src0 = dst->src[0];
-
-    const int ith = params->ith;
-    const int nth = params->nth;
-
-    const int  n_past  = ((int32_t *) dst->op_params)[0];
-    const bool inplace = src0->data == dst->data;
-
-    GGML_ASSERT(n_past >= 0);
-
-    if (!inplace) {
-        if (ith == 0) {
-            // memcpy needs to be synchronized across threads to avoid race conditions.
-            // => do it in INIT phase
-            GGML_ASSERT(ggml_nelements(dst) == ggml_nelements(src0));
-            GGML_ASSERT(ggml_is_contiguous(dst) && ggml_is_contiguous(src0));
-            memcpy(
-                ((char *)  dst->data),
-                ((char *) src0->data),
-                ggml_nbytes(dst));
-        }
-        ggml_barrier(params->threadpool);
-    }
-
-    // TODO: handle transposed/permuted matrices
-
-    const int n  = ggml_nrows(src0);
-    const int nc = src0->ne[0];
-    const int nr = src0->ne[1];
-    const int nz = n/nr;
-
-    GGML_ASSERT( dst->nb[0] == sizeof(float));
-    GGML_ASSERT(src0->nb[0] == sizeof(float));
-
-    for (int k = 0; k < nz; k++) {
-        for (int j = ith; j < nr; j += nth) {
-            for (int i = n_past; i < nc; i++) {
-                if (i > n_past + j) {
-                    *(float *)((char *) dst->data + k*dst->nb[2] + j*dst->nb[1] + i*dst->nb[0]) = value;
-                }
-            }
-        }
-    }
-}
-
-void ggml_compute_forward_diag_mask_inf(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-
-    const ggml_tensor * src0 = dst->src[0];
-
-    switch (src0->type) {
-        case GGML_TYPE_F32:
-            {
-                ggml_compute_forward_diag_mask_f32(params, dst, -INFINITY);
-            } break;
-        default:
-            {
-                GGML_ABORT("fatal error");
-            }
-    }
-}
-
-void ggml_compute_forward_diag_mask_zero(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-
-    const ggml_tensor * src0 = dst->src[0];
-
-    switch (src0->type) {
-        case GGML_TYPE_F32:
-            {
-                ggml_compute_forward_diag_mask_f32(params, dst, 0);
-            } break;
-        default:
-            {
-                GGML_ABORT("fatal error");
-            }
-    }
-}
-
 // ggml_compute_forward_soft_max
 
 static void ggml_compute_forward_soft_max_f32(

ggml/src/ggml-cpu/ops.h

@@ -56,8 +56,6 @@ void ggml_compute_forward_get_rows(const struct ggml_compute_params * params, st
 void ggml_compute_forward_get_rows_back(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_set_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_diag(const struct ggml_compute_params * params, struct ggml_tensor * dst);
-void ggml_compute_forward_diag_mask_inf(const struct ggml_compute_params * params, struct ggml_tensor * dst);
-void ggml_compute_forward_diag_mask_zero(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_soft_max(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_soft_max_ext_back(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_rope(const struct ggml_compute_params * params, struct ggml_tensor * dst);

ggml/src/ggml-cpu/simd-gemm.h

@@ -9,7 +9,7 @@
 
 // TODO: untested on avx512
 // These are in units of GGML_F32_EPR
-#if defined(__AVX512F__) || defined (__ARM_NEON__)
+#if defined(__AVX512F__) || defined(__ARM_NEON)
     static constexpr int GEMM_RM = 4;
     static constexpr int GEMM_RN = 4; // 16+4+1 = 25/32
 #elif defined(__AVX2__) || defined(__AVX__)

ggml/src/ggml-cpu/spacemit/ime.cpp

@@ -1154,9 +1154,17 @@ class tensor_traits_common : public tensor_traits_base {
                     ggml_compute_forward_concat(params, op);
                 }
                 return true;
-            // TODO For GGML_OP_GATED_DELTA_NET
-            // case GGML_OP_GATED_DELTA_NET:
-            //     return true;
+            case GGML_OP_GATED_DELTA_NET:
+                if (params->use_ref) {
+                    return false;
+                }
+                switch (op->src[0]->type) {
+                    case GGML_TYPE_F32:
+                        spacemit_kernels::rvv::forward_gated_delta_net_f32(params, op);
+                        return true;
+                    default:
+                        GGML_ABORT("fatal error");
+                }
             default:
                 break;
         }
@@ -1315,11 +1323,10 @@ static const ggml::cpu::tensor_traits * ggml_riscv64_spacemit_get_optimal_repack
         case GGML_TYPE_Q4_1:
             {
 #if defined(RISCV64_SPACEMIT_IME2)
-                // TODO
-                // if (cur->ne[1] % 32 == 0 && cur->ne[0] % 256 == 0 &&
-                //     (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
-                //     return &ggml::cpu::riscv64_spacemit::q4_1_32x256_q8_0;
-                // }
+                if (cur->ne[1] % 32 == 0 && cur->ne[0] % 256 == 0 &&
+                    (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
+                    return &ggml::cpu::riscv64_spacemit::q4_1_32x256_q8_0;
+                }
 
                 if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
                     return &ggml::cpu::riscv64_spacemit::q4_1_32x32_q8_0;
@@ -1661,7 +1668,7 @@ class extra_buffer_type : ggml::cpu::extra_buffer_type {
             case GGML_OP_SUM_ROWS:
             case GGML_OP_GET_ROWS:
             case GGML_OP_CONCAT:
-                // case GGML_OP_GATED_DELTA_NET:
+            case GGML_OP_GATED_DELTA_NET:
                 return (ggml::cpu::tensor_traits *) (&ggml::cpu::riscv64_spacemit::rvv_impl);
             default:
                 // GGML_ABORT("fatal error");

ggml/src/ggml-cpu/spacemit/ime2_kernels.cpp

@@ -2999,8 +2999,107 @@ void gemm_kernel_i8i4_hp_m1(size_t          blk_len,
                   "v24", "v25", "v26", "v27", "v28", "v29", "v30", "v31", "fa0", "fa1", "ft0", "ft1");
         }
     } else {
-        // TODO: support quant_b_zp for i8i4 hp kernel
-        GGML_ABORT("gemm_kernel_i8i4_hp_m1 with quant_b_zp is not supported yet");
+        const _Float16 hp_scale_0125 = 0.125;
+        for (size_t ni = 0; ni < count_n; ni += 32) {
+            const size_t nb_real = std::min<size_t>(NB_COLS, count_n - ni);
+            if (nb_real != NB_COLS) {
+                break;
+            }
+
+            uint8_t * b_data = (uint8_t *) quant_b_data + (ni / NB_COLS) * b_tile_stride;
+            int8_t *  a_data = (int8_t *) quant_a_ptr;
+            float *   dst_c  = c_ptr + ni;
+
+            asm volatile(
+                "vsetvli        t0, x0, e16, m1         \n\t"
+                "vxor.vv        v31, v31, v31           \n\t"
+                "mv             t4, %[BK]               \n\t"
+                "li             t0, 0x4c00              \n\t"
+                "fmv.h.x        fa0, t0                 \n\t"
+
+                ".align 4                               \n\t"
+                "BLK_LOOP%=:                            \n\t"
+                "li             t1, 4608                \n\t"
+                "add            s6, %[B], t1            \n\t"
+                "li             t5, 8                   \n\t"
+                "addi           t6, %[A], 288           \n\t"
+                "flh            ft1, (t6)               \n\t"
+                "addi           t6, %[A], 272           \n\t"
+
+                "vsetvli        t0, x0, e16, m1         \n\t"
+                "vxor.vv        v16, v18, v18           \n\t"
+                "vxor.vv        v17, v18, v18           \n\t"
+                "vxor.vv        v18, v18, v18           \n\t"
+                "vxor.vv        v19, v18, v18           \n\t"
+
+                "INNER_BLK_LOOP%=:                      \n\t"
+                "flh            fa1, (t6)               \n\t"
+                "addi           t6, t6, 2               \n\t"
+                "flh            ft0, (%[A])             \n\t"
+                "addi           %[A], %[A], 2           \n\t"
+                "vsetvli        t0, x0, e8, mf4         \n\t"
+                "vle8.v         v3, (%[A])              \n\t"
+                "addi           %[A], %[A], 32          \n\t"
+
+                "vsetvli        t0, x0, e16, mf2        \n\t"
+                "vle16.v        v8, (%[B])              \n\t"
+                "addi           %[B], %[B], 64          \n\t"
+                "vl4r.v         v4, (%[B])              \n\t"
+                "addi           %[B], %[B], 512         \n\t"
+                "vfmul.vf       v8, v8, ft0             \n\t"
+                "vfmul.vf       v9, v8, fa0             \n\t"
+                "vsetvli        t0, x0, e8, mf4         \n\t"
+                "vle8.v         v11, (s6)               \n\t"
+                "addi           s6, s6, 32              \n\t"
+                "vfwcvt.f.xu.v  v13, v11                \n\t"
+                "fmul.h         ft2, fa1, %[HP0125]     \n\t"
+                "vsetvli        t0, x0, e16, mf2        \n\t"
+                "vfmul.vv       v10, v13, v8            \n\t"
+                "vfmul.vf       v10, v10, ft2           \n\t"
+                "vfwmacc.vf     v31, ft1, v10           \n\t"
+
+                "vsetvli        t0, x0, e8, m1          \n\t"
+                "vpack.vv       v0, v8, v9, 3           \n\t"
+                "vsrl.vi        v28, v3, 4              \n\t"
+
+                "vsetvli        t0, x0, e16, m1         \n\t"
+                "vnpack4.vv     v2, v3, v3, 3           \n\t"
+                "vnpack4.vv     v3, v28, v28, 3         \n\t"
+
+                "vsetvli        t0, x0, e16, m1         \n\t"
+                "vmadotsu.hp    v16, v3, v4, v0, 4, i4  \n\t"
+                "vmadotsu.hp    v17, v3, v5, v0, 5, i4  \n\t"
+                "vmadotsu.hp    v18, v3, v6, v0, 6, i4  \n\t"
+                "vmadotsu.hp    v19, v3, v7, v0, 7, i4  \n\t"
+                "vmadotu.hp     v16, v2, v4, v0, 0, i4  \n\t"
+                "vmadotu.hp     v17, v2, v5, v0, 1, i4  \n\t"
+                "vmadotu.hp     v18, v2, v6, v0, 2, i4  \n\t"
+                "vmadotu.hp     v19, v2, v7, v0, 3, i4  \n\t"
+
+                "addi           t5, t5, -1              \n\t"
+                "bgtz           t5, INNER_BLK_LOOP%=    \n\t"
+
+                "vpack.vv       v8, v16, v17, 1         \n\t"
+                "vpack.vv       v12, v18, v19, 1        \n\t"
+                "vpack.vv       v20, v8, v12, 2         \n\t"
+
+                "vsetvli        t0, x0, e16, mf2        \n\t"
+                "addi           t4, t4, -1              \n\t"
+                "vfwmacc.vf     v31, ft1, v20           \n\t"
+
+                "addi           %[A], t6, 2             \n\t"
+                "mv             %[B], s6                \n\t"
+
+                "bgtz           t4, BLK_LOOP%=          \n\t"
+
+                "vsetvli        t0, x0, e32, m1         \n\t"
+                "vse32.v        v31, (%[DST])           \n\t"
+                : [A] "+r"(a_data), [B] "+r"(b_data)
+                : [DST] "r"(dst_c), [BK] "r"(k_blks), [HP0125] "f"(hp_scale_0125)
+                : "t0", "t1", "t2", "t3", "t4", "t5", "t6", "s6", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9",
+                  "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18", "v19", "v20", "v21", "v22", "v23",
+                  "v24", "v25", "v26", "v27", "v28", "v29", "v30", "v31", "fa0", "fa1", "ft0", "ft1", "ft2");
+        }
     }
 }
 

ggml/src/ggml-cpu/spacemit/rvv_kernels.cpp

@@ -3175,4 +3175,199 @@ template void forward_get_rows<int16_t>(ggml_compute_params * params, ggml_tenso
 template void forward_concat<int32_t>(ggml_compute_params * params, ggml_tensor * op);
 template void forward_concat<int16_t>(ggml_compute_params * params, ggml_tensor * op);
 
+static inline float rvv_dot_f32(const float * x, const float * y, int64_t n) {
+    size_t       vl   = __riscv_vsetvlmax_e32m4();
+    vfloat32m4_t vsum = __riscv_vfmv_v_f_f32m4(0.0f, vl);
+    for (int64_t i = 0; i < n; i += vl) {
+        vl                 = __riscv_vsetvl_e32m4(n - i);
+        vfloat32m4_t ax    = __riscv_vle32_v_f32m4(x + i, vl);
+        vfloat32m4_t ay    = __riscv_vle32_v_f32m4(y + i, vl);
+        vsum               = __riscv_vfmacc_vv_f32m4_tu(vsum, ax, ay, vl);
+    }
+    vl              = __riscv_vsetvlmax_e32m4();
+    vfloat32m1_t vs = __riscv_vfmv_v_f_f32m1(0.0f, 1);
+    vs              = __riscv_vfredusum_vs_f32m4_f32m1(vsum, vs, vl);
+    return __riscv_vfmv_f_s_f32m1_f32(vs);
+}
+
+static void forward_gated_delta_net_one_chunk(
+    const ggml_compute_params * params,
+    ggml_tensor * dst,
+    int64_t ir0,
+    int64_t ir1) {
+
+    ggml_tensor * src_q     = dst->src[0];
+    ggml_tensor * src_k     = dst->src[1];
+    ggml_tensor * src_v     = dst->src[2];
+    ggml_tensor * src_g     = dst->src[3];
+    ggml_tensor * src_beta  = dst->src[4];
+    ggml_tensor * src_state = dst->src[5];
+
+    const int64_t S_v      = src_v->ne[0];
+    const int64_t H        = src_v->ne[1];
+    const int64_t n_tokens = src_v->ne[2];
+    const int64_t n_seqs   = src_v->ne[3];
+
+    GGML_ASSERT(ggml_is_contiguous_rows(src_q));
+    GGML_ASSERT(ggml_is_contiguous_rows(src_k));
+    GGML_ASSERT(ggml_is_contiguous_rows(src_v));
+    GGML_ASSERT(ggml_is_contiguous(src_g));
+    GGML_ASSERT(ggml_is_contiguous(src_beta));
+    GGML_ASSERT(ggml_is_contiguous(src_state));
+
+    GGML_ASSERT(src_g->ne[0] == 1 || src_g->ne[0] == S_v);
+    GGML_ASSERT(src_beta->ne[0] == 1);
+
+    GGML_TENSOR_LOCALS(int64_t, neq, src_q, ne);
+    GGML_TENSOR_LOCALS(size_t,  nbq, src_q, nb);
+    GGML_TENSOR_LOCALS(int64_t, nek, src_k, ne);
+    GGML_TENSOR_LOCALS(size_t,  nbk, src_k, nb);
+    GGML_TENSOR_LOCALS(int64_t, nev, src_v, ne);
+    GGML_TENSOR_LOCALS(size_t,  nbv, src_v, nb);
+    GGML_TENSOR_LOCALS(int64_t, neg, src_g, ne);
+    GGML_TENSOR_LOCALS(size_t,  nbg, src_g, nb);
+    GGML_TENSOR_LOCALS(size_t,  nbb, src_beta, nb);
+
+    const bool kda = (neg0 == S_v);
+
+    const int64_t K = ggml_get_op_params_i32(dst, 0);
+    GGML_ASSERT(K >= 1);
+    const int64_t state_seq_stride = src_state->nb[3] / sizeof(float);
+
+    const int64_t per_thread = S_v + (K > 1 ? S_v * S_v : 0);
+    const int ith = params->ith;
+
+    float * delta      = (float *) params->wdata + ith * per_thread + CACHE_LINE_SIZE_F32;
+    float * state_work = K > 1 ? (delta + S_v) : nullptr;
+
+    const int64_t attn_score_elems    = S_v * H * n_tokens * n_seqs;
+    const int64_t state_size_per_snap = S_v * S_v * H * n_seqs;
+    float * attn_out_base  = (float *) dst->data;
+    float * state_out_base = (float *) dst->data + attn_score_elems;
+
+    const float * state_in_base = (const float *) src_state->data;
+
+    const int64_t rq3 = nev3 / neq3;
+    const int64_t rk3 = nev3 / nek3;
+
+    const float scale = 1.0f / sqrtf((float) S_v);
+
+    for (int64_t ir = ir0; ir < ir1; ++ir) {
+        const int64_t iv1 = ir % H;
+        const int64_t iv3 = ir / H;
+
+        const int64_t iq1 = iv1 % neq1;
+        const int64_t ik1 = iv1 % nek1;
+
+        const int64_t iq3 = iv3 / rq3;
+        const int64_t ik3 = iv3 / rk3;
+
+        float * s_out = (K > 1)
+            ? state_work
+            : state_out_base + (iv3 * H + iv1) * S_v * S_v;
+
+        const float * s_in = state_in_base + iv3 * state_seq_stride + iv1 * S_v * S_v;
+        memcpy(s_out, s_in, S_v * S_v * sizeof(float));
+
+        float * attn_data = attn_out_base + (iv3 * n_tokens * H + iv1) * S_v;
+
+        for (int64_t t = 0; t < n_tokens; t++) {
+            const float * q_d = (const float *) ((const char *) src_q->data + iq3 * nbq3 + t * nbq2 + iq1 * nbq1);
+            const float * k_d = (const float *) ((const char *) src_k->data + ik3 * nbk3 + t * nbk2 + ik1 * nbk1);
+            const float * v_d = (const float *) ((const char *) src_v->data + iv3 * nbv3 + t * nbv2 + iv1 * nbv1);
+
+            const float beta_val = *(const float *) ((const char *) src_beta->data + iv3 * nbb3 + t * nbb2 + iv1 * nbb1);
+            const float * g_d    = (const float *) ((const char *) src_g->data + iv3 * nbg3 + t * nbg2 + iv1 * nbg1);
+
+            if (kda) {
+                for (int64_t i = 0; i < S_v; ++i) {
+                    delta[i] = expf(g_d[i]);
+                }
+                for (int64_t j = 0; j < S_v; ++j) {
+                    float * row = s_out + j * S_v;
+                    const float * d = delta;
+                    int64_t remaining = S_v;
+                    while (remaining > 0) {
+                        const size_t vl = __riscv_vsetvl_e32m4(remaining);
+                        vfloat32m4_t vr = __riscv_vle32_v_f32m4(row, vl);
+                        vfloat32m4_t vd = __riscv_vle32_v_f32m4(d, vl);
+                        vr = __riscv_vfmul_vv_f32m4(vr, vd, vl);
+                        __riscv_vse32_v_f32m4(row, vr, vl);
+                        row += vl;
+                        d += vl;
+                        remaining -= vl;
+                    }
+                }
+            } else {
+                const float gate = expf(g_d[0]);
+                float * p = s_out;
+                int64_t remaining = S_v * S_v;
+                while (remaining > 0) {
+                    const size_t vl = __riscv_vsetvl_e32m4(remaining);
+                    vfloat32m4_t v = __riscv_vle32_v_f32m4(p, vl);
+                    v = __riscv_vfmul_vf_f32m4(v, gate, vl);
+                    __riscv_vse32_v_f32m4(p, v, vl);
+                    p += vl;
+                    remaining -= vl;
+                }
+            }
+
+            for (int64_t j = 0; j < S_v; ++j) {
+                const float sum = rvv_dot_f32(s_out + j * S_v, k_d, S_v);
+                delta[j] = (v_d[j] - sum) * beta_val;
+            }
+
+            for (int64_t j = 0; j < S_v; ++j) {
+                float * row = s_out + j * S_v;
+                const float * k = k_d;
+                const float dj = delta[j];
+                int64_t remaining = S_v;
+                while (remaining > 0) {
+                    const size_t vl = __riscv_vsetvl_e32m4(remaining);
+                    vfloat32m4_t vr = __riscv_vle32_v_f32m4(row, vl);
+                    vfloat32m4_t vk = __riscv_vle32_v_f32m4(k, vl);
+                    vr = __riscv_vfmacc_vf_f32m4(vr, dj, vk, vl);
+                    __riscv_vse32_v_f32m4(row, vr, vl);
+                    row += vl;
+                    k += vl;
+                    remaining -= vl;
+                }
+            }
+
+            for (int64_t j = 0; j < S_v; ++j) {
+                const float sum = rvv_dot_f32(s_out + j * S_v, q_d, S_v);
+                attn_data[j] = sum * scale;
+            }
+
+            attn_data += S_v * H;
+
+            if (K > 1) {
+                const int64_t target_slot = n_tokens - 1 - t;
+                if (target_slot >= 0 && target_slot < K) {
+                    float * curr_state_o = state_out_base + target_slot * state_size_per_snap +
+                                     (iv3 * H + iv1) * S_v * S_v;
+                    memcpy(curr_state_o, s_out, S_v * S_v * sizeof(float));
+                }
+            }
+        }
+    }
+}
+
+void forward_gated_delta_net_f32(ggml_compute_params * params, ggml_tensor * dst) {
+    ggml_tensor * src_v = dst->src[2];
+
+    int ith = params->ith;
+    int nth = params->nth;
+
+    const auto nr              = src_v->ne[1] * src_v->ne[3];
+    const auto task_per_thread = (nr + nth - 1) / nth;
+
+    const auto task_begin = ith * task_per_thread;
+    const auto task_end   = std::min((ith + 1) * task_per_thread, nr);
+
+    if (task_begin < task_end) {
+        forward_gated_delta_net_one_chunk(params, dst, task_begin, task_end);
+    }
+}
+
 }  // namespace spacemit_kernels::rvv

ggml/src/ggml-cpu/spacemit/rvv_kernels.h

@@ -62,6 +62,8 @@ void forward_rms_norm_f32(ggml_compute_params * params, ggml_tensor * op);
 
 void forward_norm_f32(ggml_compute_params * params, ggml_tensor * op);
 
+void forward_gated_delta_net_f32(ggml_compute_params * params, ggml_tensor * op);
+
 void forward_cont_with_permute(ggml_compute_params * params, ggml_tensor * op);
 
 void forward_cpy_with_permute(ggml_compute_params * params, ggml_tensor * op);

ggml/src/ggml-cuda/common.cuh

@@ -132,6 +132,14 @@ static __device__ __forceinline__ void ggml_cuda_syncwarp() {
 #endif // GGML_USE_HIP
 }
 
+static __device__ __forceinline__ void ggml_cuda_prefetch_l2(const void * p) {
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+    asm volatile("prefetch.global.L2 [%0];" :: "l"(p));
+#else
+    GGML_UNUSED(p);
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+}
+
 static __device__ __forceinline__ void ggml_cuda_pdl_sync() {
 #if defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER
     cudaGridDependencySynchronize();

ggml/src/ggml-cuda/cp-async.cuh

@@ -55,3 +55,23 @@ static __device__ __forceinline__ void cp_async_wait_all() {
     NO_DEVICE_CODE;
 #endif // CP_ASYNC_AVAILABLE
 }
+
+// Commits the copies issued since the last commit as one group.
+static __device__ __forceinline__ void cp_async_commit_group() {
+#ifdef CP_ASYNC_AVAILABLE
+    asm volatile("cp.async.commit_group;");
+#else
+    NO_DEVICE_CODE;
+#endif // CP_ASYNC_AVAILABLE
+}
+
+// Makes each thread wait until at most n of its copy groups are pending.
+// Like cp_async_wait_all, this does NOT provide any additional synchronization.
+template <int n>
+static __device__ __forceinline__ void cp_async_wait_group() {
+#ifdef CP_ASYNC_AVAILABLE
+    asm volatile("cp.async.wait_group %0;" : : "n"(n));
+#else
+    NO_DEVICE_CODE;
+#endif // CP_ASYNC_AVAILABLE
+}

ggml/src/ggml-cuda/diagmask.cu

@@ -1,40 +0,0 @@
-#include "diagmask.cuh"
-
-static __global__ void diag_mask_inf_f32(const float * x, float * dst, const int ncols, const int rows_per_channel, const int n_past) {
-    const int col = blockDim.y*blockIdx.y + threadIdx.y;
-    const int row = blockDim.x*blockIdx.x + threadIdx.x;
-
-    if (col >= ncols) {
-        return;
-    }
-
-    const int i = row*ncols + col;
-    //dst[i] = col > (n_past + row % rows_per_channel) ? -INFINITY : x[i];
-    //dst[i] = x[i] - (col > n_past + row % rows_per_channel) * INT_MAX; // equivalent within rounding error but slightly faster on GPU
-    dst[i] = x[i] - (col > n_past + row % rows_per_channel) * FLT_MAX;
-}
-
-static void diag_mask_inf_f32_cuda(const float * x, float * dst, const int ncols_x, const int nrows_x, const int rows_per_channel, const int n_past, cudaStream_t stream) {
-    const dim3 block_dims(1, CUDA_DIAG_MASK_INF_BLOCK_SIZE, 1);
-    const int block_num_x = (ncols_x + CUDA_DIAG_MASK_INF_BLOCK_SIZE - 1) / CUDA_DIAG_MASK_INF_BLOCK_SIZE;
-    const dim3 block_nums(nrows_x, block_num_x, 1);
-    diag_mask_inf_f32<<<block_nums, block_dims, 0, stream>>>(x, dst, ncols_x, rows_per_channel, n_past);
-}
-
-void ggml_cuda_op_diag_mask_inf(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
-    const ggml_tensor * src0 = dst->src[0];
-    const float * src0_d = (const float *)src0->data;
-    float * dst_d = (float *)dst->data;
-    cudaStream_t stream = ctx.stream();
-
-    GGML_ASSERT(src0->type == GGML_TYPE_F32);
-    GGML_ASSERT( dst->type == GGML_TYPE_F32);
-
-    const int64_t ne00 = src0->ne[0];
-    const int64_t ne01 = src0->ne[1];
-    const int nrows0 = ggml_nrows(src0);
-
-    const int n_past = ((int32_t *) dst->op_params)[0];
-
-    diag_mask_inf_f32_cuda(src0_d, dst_d, ne00, nrows0, ne01, n_past, stream);
-}

ggml/src/ggml-cuda/diagmask.cuh

@@ -1,5 +0,0 @@
-#include "common.cuh"
-
-#define CUDA_DIAG_MASK_INF_BLOCK_SIZE 32
-
-void ggml_cuda_op_diag_mask_inf(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

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

@@ -25,7 +25,6 @@
 #include "ggml-cuda/cpy.cuh"
 #include "ggml-cuda/cross-entropy-loss.cuh"
 #include "ggml-cuda/cumsum.cuh"
-#include "ggml-cuda/diagmask.cuh"
 #include "ggml-cuda/diag.cuh"
 #include "ggml-cuda/fattn.cuh"
 #include "ggml-cuda/fp8.cuh"
@@ -34,6 +33,7 @@
 #include "ggml-cuda/im2col.cuh"
 #include "ggml-cuda/mmf.cuh"
 #include "ggml-cuda/mmq.cuh"
+#include "ggml-cuda/mmq-few.cuh"
 #include "ggml-cuda/mmvf.cuh"
 #include "ggml-cuda/mmvq.cuh"
 #include "ggml-cuda/moe-weighted-reduction.cuh"
@@ -2431,6 +2431,10 @@ static bool ggml_cuda_should_fuse_mul_mat_q(const ggml_tensor * up, const ggml_t
     const int64_t ncols     = is_id ? y->ne[2]    : y->ne[1];
     const int64_t nchannels = is_id ? x_up->ne[2] : y->ne[2]*y->ne[3];
     const bool    use_mmvq  = is_id ? ncols <= get_mmvq_mmid_max_batch(x_up->type, cc) : ggml_cuda_should_use_mmvq(x_up->type, cc, ncols);
+    // do not fuse the matmuls that few-row MMQ takes in ggml_cuda_mul_mat
+    if (!is_id && ggml_cuda_should_use_mmq_few(x_up, y, cc)) {
+        return false;
+    }
     if (use_mmvq || !ggml_cuda_should_use_mmq(x_up->type, cc, ncols, /*n_experts =*/ is_id ? x_up->ne[2] : 0)) {
         return false;
     }
@@ -2483,10 +2487,16 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor
         ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst);
         return;
     }
-    if (ggml_cuda_should_use_mmvq(src0->type, cc, ne11)) {
+    // from a per-type number of src1 columns up to 24, few-row MMQ is faster than MMVQ and MMQ
+    const bool use_mmq_few = ggml_cuda_should_use_mmq_few(src0, src1, cc);
+    if (!use_mmq_few && ggml_cuda_should_use_mmvq(src0->type, cc, ne11)) {
         ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst);
         return;
     }
+    if (use_mmq_few) {
+        ggml_cuda_mul_mat_q_few(ctx, src0, src1, dst);
+        return;
+    }
     if (ggml_cuda_should_use_mmq(src0->type, cc, ne11, /*n_experts =*/ 0)) {
         ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst);
         return;
@@ -2731,7 +2741,6 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
             ggml_cuda_dup(ctx, dst);
             break;
         case GGML_OP_ADD:
-        case GGML_OP_ADD1: // TODO: more efficient implementation
             ggml_cuda_op_add(ctx, dst);
             break;
         case GGML_OP_ADD_ID:
@@ -2926,9 +2935,6 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
         case GGML_OP_DIAG:
             ggml_cuda_op_diag(ctx, dst);
             break;
-        case GGML_OP_DIAG_MASK_INF:
-            ggml_cuda_op_diag_mask_inf(ctx, dst);
-            break;
         case GGML_OP_SOFT_MAX:
             ggml_cuda_op_soft_max(ctx, dst);
             break;
@@ -5009,6 +5015,23 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
                 return 1;
             }
 
+            if (!with_bias && src0->type == GGML_TYPE_NVFP4 && src1->type == GGML_TYPE_F32 && mm_node->type == GGML_TYPE_F32) {
+                const int cc = ggml_cuda_info().devices[cuda_ctx->device].cc;
+                const bool is_id = op == GGML_OP_MUL_MAT_ID;
+                const int64_t ncols = is_id ? src1->ne[2] : src1->ne[1];
+                const bool use_mmvq = is_id ? ncols <= get_mmvq_mmid_max_batch(src0->type, cc) :
+                                            ggml_cuda_should_use_mmvq(src0->type, cc, ncols);
+                const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE &&
+                                               ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) &&
+                                               src0->view_src;
+                if (!use_mmvq && !bad_padding_clear &&
+                        ggml_cuda_should_use_mmq(src0->type, cc, ncols, /*n_experts =*/ is_id ? src0->ne[2] : 0)) {
+                    fusion_data.prec_src1 = ggml_prec_get_src(mm_node, 1);
+                    ggml_cuda_mul_mat_q(*cuda_ctx, src0, src1, ids, out_node, &fusion_data);
+                    return n_ops - 1;
+                }
+            }
+
             if (ggml_cuda_should_fuse_mul_mat_vec_q(mm_node)) {
                 ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, out_node, &fusion_data);
                 fused_mul_mat_vec = true;
@@ -6512,7 +6535,6 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
         case GGML_OP_PERMUTE:
         case GGML_OP_TRANSPOSE:
         case GGML_OP_ADD_ID:
-        case GGML_OP_ADD1:
         case GGML_OP_SQR:
         case GGML_OP_SQRT:
         case GGML_OP_SIN:
@@ -6559,8 +6581,6 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
         }
         case GGML_OP_CONT:
             return true;
-        case GGML_OP_DIAG_MASK_INF:
-            return true;
         case GGML_OP_SOFT_MAX:
             return true;
         case GGML_OP_SOFT_MAX_BACK: {

ggml/src/ggml-cuda/mmq-few.cu

@@ -0,0 +1,876 @@
+#include "mmq-few.cuh"
+#include "cp-async.cuh"
+#include "mma.cuh"
+#include "quantize.cuh"
+#include "vecdotq.cuh"
+
+#include <type_traits>
+
+// NVIDIA only: the kernel uses the NVIDIA int8 MMA fragment layout and cp.async.
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+
+using namespace ggml_cuda_mma;
+
+// Few-row MMQ: int8 MMA mat-mul for 2-24 src1 columns. A CUDA block has 16 src0 rows (32 for 2 column groups on large matrices).
+// Its warps split K in units of 256 values and copy raw src0 rows into their own cp.async ring, src1 is quantized to the MMA fragment layout.
+
+static constexpr int MMQ_FEW_MAX_NCG = 3;  // max. number of column groups
+static constexpr int MMQ_FEW_NROWS   = 16; // src0 rows of one MMA tile
+static constexpr int MMQ_FEW_NWARPS  = 4;
+
+template <ggml_type type>
+static constexpr __host__ __device__ int mmq_few_unit_size_of() {
+    return QK_K/ggml_cuda_type_traits<type>::qk*ggml_cuda_type_traits<type>::bs;
+}
+
+// src0 bytes per 256 values, 0 for the types without a few-row decoder.
+static constexpr __host__ __device__ int mmq_few_unit_size(const ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_Q8_0:   return mmq_few_unit_size_of<GGML_TYPE_Q8_0>();
+        case GGML_TYPE_Q3_K:   return mmq_few_unit_size_of<GGML_TYPE_Q3_K>();
+        case GGML_TYPE_Q4_K:   return mmq_few_unit_size_of<GGML_TYPE_Q4_K>();
+        case GGML_TYPE_Q5_K:   return mmq_few_unit_size_of<GGML_TYPE_Q5_K>();
+        case GGML_TYPE_Q6_K:   return mmq_few_unit_size_of<GGML_TYPE_Q6_K>();
+        case GGML_TYPE_IQ3_S:  return mmq_few_unit_size_of<GGML_TYPE_IQ3_S>();
+        case GGML_TYPE_IQ4_NL: return mmq_few_unit_size_of<GGML_TYPE_IQ4_NL>();
+        case GGML_TYPE_IQ4_XS: return mmq_few_unit_size_of<GGML_TYPE_IQ4_XS>();
+        default:               return 0;
+    }
+}
+
+// Q4_K and Q5_K also need the src1 sums for their mins.
+static constexpr __host__ __device__ mmq_q8_1_ds_layout mmq_few_ds_layout(const ggml_type type) {
+    return type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K ? MMQ_Q8_1_DS_LAYOUT_DS4 : MMQ_Q8_1_DS_LAYOUT_D4;
+}
+
+// Units of 256 values per ring stage: 1 for Q6_K, 3 column groups (fewer free registers) and 32-row tiles (2 blocks per SM), else 2.
+static constexpr __host__ __device__ int mmq_few_nunits_stage(const ggml_type type, const int ncg, const int nmt) {
+    return type == GGML_TYPE_Q6_K || ncg > 2 || nmt > 1 ? 1 : 2;
+}
+
+// A third ring stage only makes Q6_K faster, the other types get slower on large matrices.
+static constexpr __host__ __device__ int mmq_few_nstages(const ggml_type type) {
+    return type == GGML_TYPE_Q6_K ? 3 : 2;
+}
+
+// The src0 data of a ring stage does not fill whole 16 byte chunks, it is copied from the preceding 16 byte boundary.
+static constexpr __host__ __device__ bool mmq_few_unaligned(const ggml_type type, const int nunits_stage) {
+    return (nunits_stage*mmq_few_unit_size(type)) % 16 != 0;
+}
+
+// for unaligned data the max. over the 8 possible start offsets
+static constexpr __host__ __device__ int mmq_few_nchunks(const ggml_type type, const int nunits_stage) {
+    const int nbytes = nunits_stage*mmq_few_unit_size(type);
+    return mmq_few_unaligned(type, nunits_stage) ? (14 + nbytes + 15)/16 : nbytes/16;
+}
+
+// Unaligned rows are padded to an odd number of 16 byte chunks against bank conflicts, for aligned rows the padding costs more than it saves.
+static constexpr __host__ __device__ int mmq_few_sram_stride(const ggml_type type, const int nunits_stage) {
+    const int nchunks = mmq_few_nchunks(type, nunits_stage);
+    return mmq_few_unaligned(type, nunits_stage) && nchunks % 2 == 0 ? 16*(nchunks + 1) : 16*nchunks;
+}
+
+static constexpr __host__ __device__ int mmq_few_nbytes_shared(const ggml_type type, const int ncg, const int nmt) {
+    return MMQ_FEW_NWARPS*mmq_few_nstages(type)*nmt*MMQ_FEW_NROWS*mmq_few_sram_stride(type, mmq_few_nunits_stage(type, ncg, nmt));
+}
+
+// With Q6_K and Q8_0 two CUDA blocks with 32-row tiles would not fit into an SM.
+static constexpr __host__ __device__ bool mmq_few_tall_supported(const ggml_type type) {
+    return type == GGML_TYPE_Q3_K || type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K ||
+           type == GGML_TYPE_IQ3_S || type == GGML_TYPE_IQ4_NL || type == GGML_TYPE_IQ4_XS;
+}
+
+#if defined(CP_ASYNC_AVAILABLE)
+
+// The src1 fragments of a lane for one unit of 256 values, see quantize.cuh.
+struct mmq_few_y_frag {
+    int4 qs[4];
+    int4 ds[4];
+};
+
+static __device__ __forceinline__ void mmq_few_load_y(mmq_few_y_frag & yf, const int4 * y) {
+#pragma unroll
+    for (int j = 0; j < 4; ++j) {
+        yf.qs[j] = y[j*WARP_SIZE + threadIdx.x];
+    }
+#pragma unroll
+    for (int j = 0; j < 4; ++j) {
+        yf.ds[j] = y[4*WARP_SIZE + 4*(threadIdx.x % 4) + j];
+    }
+}
+
+static __device__ __forceinline__ int2 mmq_few_get_b(const mmq_few_y_frag & yf, const int s) {
+    const int4 q = yf.qs[s/2];
+    return s % 2 == 0 ? make_int2(q.x, q.y) : make_int2(q.z, q.w);
+}
+
+template <mmq_q8_1_ds_layout ds_layout>
+static __device__ __forceinline__ float2 mmq_few_get_d(const mmq_few_y_frag & yf, const int s) {
+    if constexpr (ds_layout == MMQ_Q8_1_DS_LAYOUT_D4) {
+        const float * d = (const float *) &yf.ds[s/2];
+        return make_float2(d[2*(s % 2) + 0], d[2*(s % 2) + 1]);
+    } else {
+        const half2 * ds = (const half2 *) &yf.ds[s/2];
+        return make_float2(__low2float(ds[2*(s % 2) + 0]), __low2float(ds[2*(s % 2) + 1]));
+    }
+}
+
+static __device__ __forceinline__ float2 mmq_few_get_sum(const mmq_few_y_frag & yf, const int s) {
+    const half2 * ds = (const half2 *) &yf.ds[s/2];
+    return make_float2(__high2float(ds[2*(s % 2) + 0]), __high2float(ds[2*(s % 2) + 1]));
+}
+
+// C.x[0], C.x[1]: row g, columns 2t and 2t+1. C.x[2], C.x[3]: row g+8.
+static __device__ __forceinline__ tile<16, 8, int> mmq_few_mma(const int (&lo)[2], const int (&hi)[2], const int2 b) {
+    tile<16, 8, int> A;
+    A.x[0] = lo[0];
+    A.x[1] = lo[1];
+    A.x[2] = hi[0];
+    A.x[3] = hi[1];
+    tile<8, 8, int> B;
+    B.x[0] = b.x;
+    B.x[1] = b.y;
+    tile<16, 8, int> C;
+    mma(C, A, B);
+    return C;
+}
+
+static __device__ __forceinline__ tile<16, 8, int> mmq_few_mma_k16(const int (&a)[2], const int b) {
+    tile<16, 4, int> A;
+    A.x[0] = a[0];
+    A.x[1] = a[1];
+    tile<8, 4, int> B;
+    B.x[0] = b;
+    tile<16, 8, int> C;
+    mma(C, A, B);
+    return C;
+}
+
+// sum_d += d * (C * ls) of sub-block s for all column groups, ls are the src0 row scales
+template <ggml_type type, int ncg, typename scale_t>
+static __device__ __forceinline__ void mmq_few_mma_add(
+        const int (&lo)[2], const int (&hi)[2], const scale_t (&ls)[2], const mmq_few_y_frag (&yf)[ncg], const int s, float (&sum_d)[ncg][4]) {
+#pragma unroll
+    for (int c = 0; c < ncg; ++c) {
+        const tile<16, 8, int> C = mmq_few_mma(lo, hi, mmq_few_get_b(yf[c], s));
+
+        const float2 d = mmq_few_get_d<mmq_few_ds_layout(type)>(yf[c], s);
+        sum_d[c][0] += d.x * float(C.x[0]*ls[0]);
+        sum_d[c][1] += d.y * float(C.x[1]*ls[0]);
+        sum_d[c][2] += d.x * float(C.x[2]*ls[1]);
+        sum_d[c][3] += d.y * float(C.x[3]*ls[1]);
+    }
+}
+
+// Same with one k=16 MMA per half of sub-block s, the int8 scales of the halves are bytes 0 and 1 of sc[r].
+template <ggml_type type, int ncg>
+static __device__ __forceinline__ void mmq_few_mma_k16_add(
+        const int (&lo)[2], const int (&hi)[2], const int (&sc)[2], const mmq_few_y_frag (&yf)[ncg], const int s, float (&sum_d)[ncg][4]) {
+    const int sc_lo[2] = {(int8_t) sc[0], (int8_t) sc[1]};
+    const int sc_hi[2] = {(int8_t) (sc[0] >> 8), (int8_t) (sc[1] >> 8)};
+#pragma unroll
+    for (int c = 0; c < ncg; ++c) {
+        const int2 b = mmq_few_get_b(yf[c], s);
+        const tile<16, 8, int> C_lo = mmq_few_mma_k16(lo, b.x);
+        const tile<16, 8, int> C_hi = mmq_few_mma_k16(hi, b.y);
+
+        const float2 d = mmq_few_get_d<mmq_few_ds_layout(type)>(yf[c], s);
+        sum_d[c][0] += d.x * float(C_lo.x[0]*sc_lo[0] + C_hi.x[0]*sc_hi[0]);
+        sum_d[c][1] += d.y * float(C_lo.x[1]*sc_lo[0] + C_hi.x[1]*sc_hi[0]);
+        sum_d[c][2] += d.x * float(C_lo.x[2]*sc_lo[1] + C_hi.x[2]*sc_hi[1]);
+        sum_d[c][3] += d.y * float(C_lo.x[3]*sc_lo[1] + C_hi.x[3]*sc_hi[1]);
+    }
+}
+
+template <int ncg>
+static __device__ __forceinline__ void mmq_few_add_scaled(float (&sum)[ncg][4], const float (&d)[2], const float (&sum_d)[ncg][4]) {
+#pragma unroll
+    for (int c = 0; c < ncg; ++c) {
+#pragma unroll
+        for (int l = 0; l < 4; ++l) {
+            sum[c][l] += d[l/2]*sum_d[c][l];
+        }
+    }
+}
+
+// vec_dot: one unit of 256 values of the 16 rows of a warp, x[0] and x[1] are the rows g and g+8 of lane (g, t).
+// The lane decodes the values [4t, 4t + 4) and [16 + 4t, 16 + 4t + 4) of each sub-block, its A fragments.
+
+template <ggml_type type, int ncg>
+static __device__ __forceinline__ void mmq_few_vec_dot_q45_K(const char * const (&x)[2], const mmq_few_y_frag (&yf)[ncg], float (&sum)[ncg][4]) {
+    using block_t = std::conditional_t<type == GGML_TYPE_Q4_K, block_q4_K, block_q5_K>;
+    const int t = threadIdx.x % 4;
+
+    int   qs[2][4][2];
+    [[maybe_unused]] int qh[2][2]; // Q5_K: bit s holds the high bit of sub-block s
+    int   sc[2][2];
+    int   mn[2][2];
+    float d[2];
+    float dmin[2];
+#pragma unroll
+    for (int r = 0; r < 2; ++r) {
+        const block_t * bx = (const block_t *) x[r];
+#pragma unroll
+        for (int j = 0; j < 4; ++j) {
+            qs[r][j][0] = get_int_b4(bx->qs + 32*j,      t);
+            qs[r][j][1] = get_int_b4(bx->qs + 32*j + 16, t);
+        }
+        if constexpr (type == GGML_TYPE_Q5_K) {
+            qh[r][0] = get_int_b4(bx->qh,      t);
+            qh[r][1] = get_int_b4(bx->qh + 16, t);
+        }
+        const int4 hdr = *(const int4 *) bx; // dm and the 12 scale bytes
+        const int scales[3] = {hdr.y, hdr.z, hdr.w};
+        sc[r][0] = unpack_scales_q45_K(scales, 0);
+        sc[r][1] = unpack_scales_q45_K(scales, 1);
+        mn[r][0] = unpack_scales_q45_K(scales, 2);
+        mn[r][1] = unpack_scales_q45_K(scales, 3);
+        const float2 dm = __half22float2(bx->dm);
+        d[r]    = dm.x;
+        dmin[r] = dm.y;
+    }
+
+    float sum_d[ncg][4] = {{0.0f}};
+    float sum_m[ncg][4] = {{0.0f}};
+#pragma unroll
+    for (int s = 0; s < 8; ++s) {
+        const int j     = s/2;
+        const int shift = 4*(s % 2);
+        int lo[2];
+        int hi[2];
+#pragma unroll
+        for (int r = 0; r < 2; ++r) {
+            lo[r] = (qs[r][j][0] >> shift) & 0x0F0F0F0F;
+            hi[r] = (qs[r][j][1] >> shift) & 0x0F0F0F0F;
+            if constexpr (type == GGML_TYPE_Q5_K) {
+                lo[r] |= ((qh[r][0] >> s) & 0x01010101) << 4;
+                hi[r] |= ((qh[r][1] >> s) & 0x01010101) << 4;
+            }
+        }
+        const int   ls[2] = {(sc[0][s/4] >> (8*(s % 4))) & 0xFF, (sc[1][s/4] >> (8*(s % 4))) & 0xFF};
+        const float m0    = (mn[0][s/4] >> (8*(s % 4))) & 0xFF;
+        const float m1    = (mn[1][s/4] >> (8*(s % 4))) & 0xFF;
+        mmq_few_mma_add<type>(lo, hi, ls, yf, s, sum_d);
+#pragma unroll
+        for (int c = 0; c < ncg; ++c) {
+            const float2 sm = mmq_few_get_sum(yf[c], s);
+            sum_m[c][0] += m0 * sm.x;
+            sum_m[c][1] += m0 * sm.y;
+            sum_m[c][2] += m1 * sm.x;
+            sum_m[c][3] += m1 * sm.y;
+        }
+    }
+#pragma unroll
+    for (int c = 0; c < ncg; ++c) {
+#pragma unroll
+        for (int l = 0; l < 4; ++l) {
+            sum[c][l] += d[l/2]*sum_d[c][l] - dmin[l/2]*sum_m[c][l];
+        }
+    }
+}
+
+// Q6_K and Q3_K have a scale per 16 values: one k=16 MMA per half of a sub-block.
+template <int ncg>
+static __device__ __forceinline__ void mmq_few_vec_dot_q6_K(const char * const (&x)[2], const mmq_few_y_frag (&yf)[ncg], float (&sum)[ncg][4]) {
+    const int t = threadIdx.x % 4;
+
+    float d[2];
+#pragma unroll
+    for (int r = 0; r < 2; ++r) {
+        d[r] = __half2float(((const block_q6_K *) x[r])->d);
+    }
+
+    float sum_d[ncg][4] = {{0.0f}};
+#pragma unroll
+    for (int n = 0; n < 2; ++n) {
+        int ql[2][2][2]; // [row][ql offset 0 or 32][lo or hi]
+        int qh[2][2];
+        int sc[2][2];
+#pragma unroll
+        for (int r = 0; r < 2; ++r) {
+            const block_q6_K * bx = (const block_q6_K *) x[r];
+#pragma unroll
+            for (int k = 0; k < 2; ++k) {
+                ql[r][k][0] = get_int_b2(bx->ql + 64*n + 32*k,      t);
+                ql[r][k][1] = get_int_b2(bx->ql + 64*n + 32*k + 16, t);
+            }
+            qh[r][0] = get_int_b2(bx->qh + 32*n,      t);
+            qh[r][1] = get_int_b2(bx->qh + 32*n + 16, t);
+            sc[r][0] = get_int_b2(bx->scales + 8*n, 0);
+            sc[r][1] = get_int_b2(bx->scales + 8*n, 1);
+        }
+#pragma unroll
+        for (int j = 0; j < 4; ++j) {
+            const int s = 4*n + j;
+            int lo[2];
+            int hi[2];
+#pragma unroll
+            for (int r = 0; r < 2; ++r) {
+                lo[r] = __vsub4(((ql[r][j % 2][0] >> (4*(j/2))) & 0x0F0F0F0F) | (((qh[r][0] >> (2*j)) & 0x03030303) << 4), 0x20202020);
+                hi[r] = __vsub4(((ql[r][j % 2][1] >> (4*(j/2))) & 0x0F0F0F0F) | (((qh[r][1] >> (2*j)) & 0x03030303) << 4), 0x20202020);
+            }
+            const int sc_s[2] = {sc[0][j/2] >> (16*(j % 2)), sc[1][j/2] >> (16*(j % 2))}; // scales 2s and 2s+1
+            mmq_few_mma_k16_add<GGML_TYPE_Q6_K>(lo, hi, sc_s, yf, s, sum_d);
+        }
+    }
+    mmq_few_add_scaled(sum, d, sum_d);
+}
+
+template <int ncg>
+static __device__ __forceinline__ void mmq_few_vec_dot_q3_K(const char * const (&x)[2], const mmq_few_y_frag (&yf)[ncg], float (&sum)[ncg][4]) {
+    const int t = threadIdx.x % 4;
+
+    int   hm[2][2];
+    int   sc[2][4]; // the 16 scales of a row as int8, minus 32
+    float d[2];
+#pragma unroll
+    for (int r = 0; r < 2; ++r) {
+        const block_q3_K * bx = (const block_q3_K *) x[r];
+        hm[r][0] = ~get_int_b2(bx->hmask,      t); // a set bit means that 4 is subtracted
+        hm[r][1] = ~get_int_b2(bx->hmask + 16, t);
+#pragma unroll
+        for (int ksc = 0; ksc < 4; ++ksc) {
+            sc[r][ksc] = unpack_scales_q3_K(bx->scales, ksc);
+        }
+        d[r] = __half2float(bx->d);
+    }
+
+    float sum_d[ncg][4] = {{0.0f}};
+#pragma unroll
+    for (int n = 0; n < 2; ++n) {
+        int qs[2][2];
+#pragma unroll
+        for (int r = 0; r < 2; ++r) {
+            const block_q3_K * bx = (const block_q3_K *) x[r];
+            qs[r][0] = get_int_b2(bx->qs + 32*n,      t);
+            qs[r][1] = get_int_b2(bx->qs + 32*n + 16, t);
+        }
+#pragma unroll
+        for (int j = 0; j < 4; ++j) {
+            const int s = 4*n + j;
+            int lo[2];
+            int hi[2];
+#pragma unroll
+            for (int r = 0; r < 2; ++r) {
+                lo[r] = __vsub4((qs[r][0] >> (2*j)) & 0x03030303, ((hm[r][0] >> s) & 0x01010101) << 2);
+                hi[r] = __vsub4((qs[r][1] >> (2*j)) & 0x03030303, ((hm[r][1] >> s) & 0x01010101) << 2);
+            }
+            const int sc_s[2] = {sc[0][s/2] >> (16*(s % 2)), sc[1][s/2] >> (16*(s % 2))}; // scales 2s and 2s+1
+            mmq_few_mma_k16_add<GGML_TYPE_Q3_K>(lo, hi, sc_s, yf, s, sum_d);
+        }
+    }
+    mmq_few_add_scaled(sum, d, sum_d);
+}
+
+template <int ncg>
+static __device__ __forceinline__ void mmq_few_vec_dot_iq4_xs(const char * const (&x)[2], const mmq_few_y_frag (&yf)[ncg], float (&sum)[ncg][4]) {
+    const int t = threadIdx.x % 4;
+
+    float d[2];
+    int   scales_h[2];
+    int   scales_l[2];
+#pragma unroll
+    for (int r = 0; r < 2; ++r) {
+        const block_iq4_xs * bx = (const block_iq4_xs *) x[r];
+        d[r]        = __half2float(bx->d);
+        scales_h[r] = bx->scales_h;
+        scales_l[r] = get_int_b4(bx->scales_l, 0);
+    }
+
+    float sum_d[ncg][4] = {{0.0f}};
+#pragma unroll
+    for (int s = 0; s < 8; ++s) {
+        int lo[2];
+        int hi[2];
+        int ls[2];
+#pragma unroll
+        for (int r = 0; r < 2; ++r) {
+            const block_iq4_xs * bx = (const block_iq4_xs *) x[r];
+            const int2 v = get_int_from_table_16(get_int_b4(bx->qs + 16*s, t), kvalues_iq4nl);
+            lo[r] = v.x;
+            hi[r] = v.y;
+            ls[r] = (((scales_l[r] >> (4*s)) & 0x0F) | (((scales_h[r] >> (2*s)) & 0x03) << 4)) - 32;
+        }
+        mmq_few_mma_add<GGML_TYPE_IQ4_XS>(lo, hi, ls, yf, s, sum_d);
+    }
+    mmq_few_add_scaled(sum, d, sum_d);
+}
+
+// Q8_0 and IQ4_NL have a block with its own scale per 32 values.
+template <ggml_type type, int ncg>
+static __device__ __forceinline__ void mmq_few_vec_dot_q8_0_iq4_nl(const char * const (&x)[2], const mmq_few_y_frag (&yf)[ncg], float (&sum)[ncg][4]) {
+    const int t = threadIdx.x % 4;
+
+#pragma unroll
+    for (int s = 0; s < 8; ++s) {
+        int   lo[2];
+        int   hi[2];
+        float d[2];
+#pragma unroll
+        for (int r = 0; r < 2; ++r) {
+            if constexpr (type == GGML_TYPE_Q8_0) {
+                const block_q8_0 * bx = (const block_q8_0 *) x[r] + s;
+                lo[r] = get_int_b2(bx->qs,      t);
+                hi[r] = get_int_b2(bx->qs + 16, t);
+                d[r]  = __half2float(bx->d);
+            } else {
+                const block_iq4_nl * bx = (const block_iq4_nl *) x[r] + s;
+                const int2 v = get_int_from_table_16(get_int_b2(bx->qs, t), kvalues_iq4nl);
+                lo[r] = v.x;
+                hi[r] = v.y;
+                d[r]  = __half2float(bx->d);
+            }
+        }
+        mmq_few_mma_add<type>(lo, hi, d, yf, s, sum);
+    }
+}
+
+template <int ncg>
+static __device__ __forceinline__ void mmq_few_vec_dot_iq3_s(const char * const (&x)[2], const mmq_few_y_frag (&yf)[ncg], float (&sum)[ncg][4]) {
+    const int t = threadIdx.x % 4;
+
+    float d[2];
+#pragma unroll
+    for (int r = 0; r < 2; ++r) {
+        d[r] = __half2float(((const block_iq3_s *) x[r])->d);
+    }
+
+    float sum_d[ncg][4] = {{0.0f}};
+#pragma unroll
+    for (int s = 0; s < 8; ++s) {
+        int lo[2];
+        int hi[2];
+        int ls[2];
+#pragma unroll
+        for (int r = 0; r < 2; ++r) {
+            const block_iq3_s * bx = (const block_iq3_s *) x[r];
+            const int qh = bx->qh[s];
+            // the values [4k, 4k + 4) of a sub-block come from grid entry k, with the signs in bits 4*(k%2) of byte k/2
+            const int k_lo = t;
+            const int k_hi = 4 + t;
+            const int grid_lo = iq3s_grid[bx->qs[8*s + k_lo] | (((qh >> k_lo) & 1) << 8)];
+            const int grid_hi = iq3s_grid[bx->qs[8*s + k_hi] | (((qh >> k_hi) & 1) << 8)];
+            const int sign_lo = (bx->signs[4*s + k_lo/2] >> (4*(k_lo % 2))) & 0x0F;
+            const int sign_hi = (bx->signs[4*s + k_hi/2] >> (4*(k_hi % 2))) & 0x0F;
+            const int mask_lo = __vcmpne4(((sign_lo & 0x03) << 7) | ((sign_lo & 0x0C) << 21), 0x00000000);
+            const int mask_hi = __vcmpne4(((sign_hi & 0x03) << 7) | ((sign_hi & 0x0C) << 21), 0x00000000);
+            lo[r] = __vsub4(grid_lo ^ mask_lo, mask_lo);
+            hi[r] = __vsub4(grid_hi ^ mask_hi, mask_hi);
+            ls[r] = 1 + 2*((bx->scales[s/2] >> (4*(s % 2))) & 0x0F);
+        }
+        mmq_few_mma_add<GGML_TYPE_IQ3_S>(lo, hi, ls, yf, s, sum_d);
+    }
+    mmq_few_add_scaled(sum, d, sum_d);
+}
+
+template <ggml_type type, int ncg>
+static __device__ __forceinline__ void mmq_few_vec_dot(const char * const (&x)[2], const mmq_few_y_frag (&yf)[ncg], float (&sum)[ncg][4]) {
+    if constexpr (type == GGML_TYPE_Q8_0 || type == GGML_TYPE_IQ4_NL) {
+        mmq_few_vec_dot_q8_0_iq4_nl<type, ncg>(x, yf, sum);
+    } else if constexpr (type == GGML_TYPE_Q3_K) {
+        mmq_few_vec_dot_q3_K<ncg>(x, yf, sum);
+    } else if constexpr (type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K) {
+        mmq_few_vec_dot_q45_K<type, ncg>(x, yf, sum);
+    } else if constexpr (type == GGML_TYPE_Q6_K) {
+        mmq_few_vec_dot_q6_K<ncg>(x, yf, sum);
+    } else if constexpr (type == GGML_TYPE_IQ3_S) {
+        mmq_few_vec_dot_iq3_s<ncg>(x, yf, sum);
+    } else if constexpr (type == GGML_TYPE_IQ4_XS) {
+        mmq_few_vec_dot_iq4_xs<ncg>(x, yf, sum);
+    } else {
+        static_assert(type == GGML_TYPE_COUNT, "unsupported type");
+    }
+}
+
+#endif // defined(CP_ASYNC_AVAILABLE)
+
+template <ggml_type type, int ncg, int nmt>
+__launch_bounds__(MMQ_FEW_NWARPS*ggml_cuda_get_physical_warp_size(), 1)
+static __global__ void mul_mat_q_few(
+        const char * x, const int4 * y_ptr, float * dst,
+        const int nunits, const int nrows_x, const int ncols_dst, const int64_t stride_row_x, const int stride_col_dst,
+        const int channel_ratio, const int64_t stride_channel_x, const int stride_channel_dst,
+        const int sample_ratio, const int64_t stride_sample_x, const int stride_sample_dst, const bool x_is_constant) {
+#if defined(CP_ASYNC_AVAILABLE)
+    constexpr int nwarps    = MMQ_FEW_NWARPS;
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    static_assert(warp_size == WARP_SIZE, "unsupported warp size");
+
+    constexpr int  unit_size    = mmq_few_unit_size(type);
+    constexpr int  nrows        = nmt*MMQ_FEW_NROWS;
+    constexpr int  nunits_stage = mmq_few_nunits_stage(type, ncg, nmt);
+    constexpr int  nstages      = mmq_few_nstages(type);
+    constexpr bool unaligned    = mmq_few_unaligned(type, nunits_stage);
+    constexpr int  nchunks      = mmq_few_nchunks(type, nunits_stage);
+    constexpr int  sram_stride  = mmq_few_sram_stride(type, nunits_stage);
+    constexpr int  nbytes_stage = nrows*sram_stride;
+    constexpr int  ncopies      = nrows*nchunks;
+    static_assert(unit_size > 0, "unsupported type");
+
+    extern __shared__ char data_mmq_few[];
+    char * ring = data_mmq_few + threadIdx.y*nstages*nbytes_stage;
+
+    const int row0        = blockIdx.x*nrows;
+    const int channel_dst = blockIdx.y;
+    const int sample_dst  = blockIdx.z;
+    const int nchannels   = gridDim.y;
+
+    x   += int64_t(sample_dst/sample_ratio)*stride_sample_x + int64_t(channel_dst/channel_ratio)*stride_channel_x;
+    dst += int64_t(sample_dst)*stride_sample_dst + int64_t(channel_dst)*stride_channel_dst;
+    const int4 * GGML_CUDA_RESTRICT y = y_ptr + (int64_t(sample_dst)*nchannels + channel_dst)*nunits*ncg*MMQ_FEW_Q8_1_INT4;
+
+    // No bounds check for the units, the gate makes nunits a multiple of nunits_stage.
+    // The extra bytes of unaligned data are inside the allocation, which ggml aligns and pads to 128 bytes.
+    auto load_stage = [&](const int stage, const int iu0) {
+        const unsigned int dst_stage = ggml_cuda_cvta_generic_to_shared(ring + stage*nbytes_stage);
+#pragma unroll
+        for (int c0 = 0; c0 < ncopies; c0 += warp_size) {
+            const int c  = c0 + threadIdx.x;
+            const int i  = c / nchunks;
+            const int ic = c % nchunks;
+            if (ncopies % warp_size == 0 || c < ncopies) {
+                const int row = min(row0 + i, nrows_x - 1);
+                const char * src = x + int64_t(row)*stride_row_x + int64_t(iu0)*unit_size;
+                if constexpr (unaligned) {
+                    const int off = (uintptr_t) src % 16;
+                    if (16*ic < off + nunits_stage*unit_size) {
+                        cp_async_cg_16<0>(dst_stage + i*sram_stride + 16*ic, src - off + 16*ic);
+                    }
+                } else {
+                    cp_async_cg_16<0>(dst_stage + i*sram_stride + 16*ic, src + 16*ic);
+                }
+            }
+        }
+    };
+
+    // offsets of the rows of this lane from their 16 byte boundary
+    int off_row[2*nmt] = {0};
+    if constexpr (unaligned) {
+#pragma unroll
+        for (int r = 0; r < 2*nmt; ++r) {
+            const int row = min(row0 + int(threadIdx.x)/4 + 8*r, nrows_x - 1);
+            off_row[r] = (uintptr_t) (x + int64_t(row)*stride_row_x) % 16;
+        }
+    }
+
+    // Only weights, which no preceding kernel writes, are read before the PDL wait.
+    if (!x_is_constant) {
+        ggml_cuda_pdl_sync();
+    }
+
+#if __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
+    // DGX Spark DRAM serves the short row pieces of the stage copies slowly, so the rows of the block are first prefetched into L2 in address order.
+    {
+        const int     tid       = threadIdx.y*warp_size + threadIdx.x;
+        const int     nrows_blk = min(nrows, nrows_x - row0);
+        const int64_t row_bytes = int64_t(nunits)*unit_size;
+        const bool    contig    = stride_row_x == row_bytes;
+        for (int ir = 0; ir < (contig ? 1 : nrows_blk); ++ir) {
+            const uintptr_t p0 = (uintptr_t) (x + int64_t(row0 + ir)*stride_row_x);
+            const uintptr_t p1 = p0 + (contig ? nrows_blk*row_bytes : row_bytes);
+            for (uintptr_t p = (p0 & ~uintptr_t(127)) + 128*tid; p < p1; p += 128*nwarps*warp_size) {
+                ggml_cuda_prefetch_l2((const void *) p);
+            }
+        }
+    }
+#endif // __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
+
+    const int iu_step = nunits_stage*nwarps;
+
+#pragma unroll
+    for (int stage = 0; stage < nstages - 1; ++stage) {
+        const int iu0 = nunits_stage*threadIdx.y + stage*iu_step;
+        if (iu0 < nunits) {
+            load_stage(stage, iu0);
+        }
+        cp_async_commit_group();
+    }
+    if (x_is_constant) {
+        ggml_cuda_pdl_sync();
+    }
+
+    auto load_y = [&](mmq_few_y_frag (&yf)[ncg], const int iu) {
+#pragma unroll
+        for (int c = 0; c < ncg; ++c) {
+            mmq_few_load_y(yf[c], y + (int64_t(iu)*ncg + c)*MMQ_FEW_Q8_1_INT4);
+        }
+    };
+
+    // The next src1 fragments are loaded right after the MMAs that use the current ones, an earlier load is slower.
+    mmq_few_y_frag yf[nunits_stage][ncg];
+    if (nunits_stage*threadIdx.y < nunits) {
+#pragma unroll
+        for (int k = 0; k < nunits_stage; ++k) {
+            load_y(yf[k], nunits_stage*threadIdx.y + k);
+        }
+    }
+
+    float sum[nmt][ncg][4] = {{{0.0f}}};
+
+    for (int iter = 0, iu0 = nunits_stage*threadIdx.y; iu0 < nunits; ++iter, iu0 += iu_step) {
+        const int iu_load = iu0 + (nstages - 1)*iu_step;
+        if (iu_load < nunits) {
+            load_stage((iter + nstages - 1) % nstages, iu_load);
+        }
+        cp_async_commit_group();
+
+        const bool has_next = iu0 + iu_step < nunits;
+
+        cp_async_wait_group<nstages - 1>();
+        ggml_cuda_syncwarp();
+
+        const char * xs = ring + (iter % nstages)*nbytes_stage;
+        const int off_stage = unaligned ? (iu0*unit_size) % 16 : 0;
+#pragma unroll
+        for (int k = 0; k < nunits_stage; ++k) {
+#pragma unroll
+            for (int m = 0; m < nmt; ++m) {
+                const char * const xr[2] = {
+                    xs + (MMQ_FEW_NROWS*m + threadIdx.x/4 + 0)*sram_stride + (off_row[2*m + 0] + off_stage) % 16 + k*unit_size,
+                    xs + (MMQ_FEW_NROWS*m + threadIdx.x/4 + 8)*sram_stride + (off_row[2*m + 1] + off_stage) % 16 + k*unit_size,
+                };
+                mmq_few_vec_dot<type, ncg>(xr, yf[k], sum[m]);
+            }
+            if (has_next) {
+                load_y(yf[k], iu0 + iu_step + k);
+            }
+        }
+
+        ggml_cuda_syncwarp(); // the stage is overwritten in the next iteration
+    }
+    cp_async_wait_all();
+
+    ggml_cuda_pdl_lc();
+
+    // reduce the sums over the warps as one flat array
+    constexpr int nsum = nmt*ncg*4;
+    float * sum_flat = &sum[0][0][0];
+
+    float * buf_sum = (float *) data_mmq_few;
+    __syncthreads();
+    if (threadIdx.y > 0) {
+#pragma unroll
+        for (int i = 0; i < nsum; ++i) {
+            buf_sum[((threadIdx.y - 1)*nsum + i)*warp_size + threadIdx.x] = sum_flat[i];
+        }
+    }
+    __syncthreads();
+    if (threadIdx.y > 0) {
+        return;
+    }
+#pragma unroll
+    for (int w = 0; w < nwarps - 1; ++w) {
+#pragma unroll
+        for (int i = 0; i < nsum; ++i) {
+            sum_flat[i] += buf_sum[(w*nsum + i)*warp_size + threadIdx.x];
+        }
+    }
+
+#pragma unroll
+    for (int m = 0; m < nmt; ++m) {
+#pragma unroll
+        for (int c = 0; c < ncg; ++c) {
+#pragma unroll
+            for (int l = 0; l < 4; ++l) {
+                const int i = row0 + MMQ_FEW_NROWS*m + tile<16, 8, int>::get_i(l);
+                const int j = c*MMQ_FEW_NCOLS + tile<16, 8, int>::get_j(l);
+                if (i < nrows_x && j < ncols_dst) {
+                    dst[int64_t(j)*stride_col_dst + i] = sum[m][c][l];
+                }
+            }
+        }
+    }
+#else
+    GGML_UNUSED_VARS(x, y_ptr, dst, nunits, nrows_x, ncols_dst, stride_row_x, stride_col_dst,
+        channel_ratio, stride_channel_x, stride_channel_dst,
+        sample_ratio, stride_sample_x, stride_sample_dst, x_is_constant);
+    NO_DEVICE_CODE;
+#endif // defined(CP_ASYNC_AVAILABLE)
+}
+
+// ggml_is_constant (static in ggml.c) on the tensor that a view points to
+static bool mmq_few_is_constant(const ggml_tensor * src0) {
+    const ggml_tensor * x_root = src0->view_src ? src0->view_src : src0;
+    return x_root->buffer && ggml_backend_buffer_get_usage(x_root->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS &&
+        !(x_root->flags & GGML_TENSOR_FLAG_PARAM);
+}
+
+template <ggml_type type, int ncg, int nmt>
+static void launch_mul_mat_q_few(const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst,
+        const int4 * src1_q, cudaStream_t stream) {
+    GGML_TENSOR_BINARY_OP_LOCALS;
+
+    const int nunits = ne00 / QK_K;
+    GGML_ASSERT(nunits % mmq_few_nunits_stage(type, ncg, nmt) == 0);
+
+    const int warp_size     = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
+    const int nbytes_shared = mmq_few_nbytes_shared(type, ncg, nmt);
+    CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q_few<type, ncg, nmt>), nbytes_shared);
+
+    const bool x_is_constant = mmq_few_is_constant(src0);
+
+    const dim3 block_nums((ne01 + nmt*MMQ_FEW_NROWS - 1) / (nmt*MMQ_FEW_NROWS), ne12, ne13);
+    const dim3 block_dims(warp_size, MMQ_FEW_NWARPS, 1);
+    const ggml_cuda_kernel_launch_params params(block_nums, block_dims, nbytes_shared, stream);
+    ggml_cuda_kernel_launch(mul_mat_q_few<type, ncg, nmt>, params,
+        (const char *) src0->data, src1_q, (float *) dst->data,
+        nunits, (int) ne01, (int) ne1, (int64_t) nb01, (int) (nb1/sizeof(float)),
+        (int) (ne12/ne02), (int64_t) nb02, (int) (nb2/sizeof(float)),
+        (int) (ne13/ne03), (int64_t) nb03, (int) (nb3/sizeof(float)), x_is_constant);
+}
+
+static int mmq_few_ncg(const int64_t ncols) {
+    return (ncols + MMQ_FEW_NCOLS - 1) / MMQ_FEW_NCOLS;
+}
+
+// With 2 column groups, 32-row tiles halve the src1 loads per src0 byte. With 3 column groups they spill registers.
+static int mmq_few_nmt(const ggml_type type, const int ncg, const int64_t nrows_x) {
+    const int nsm = ggml_cuda_info().devices[ggml_cuda_get_device()].nsm;
+    return mmq_few_tall_supported(type) && ncg == 2 && nrows_x >= int64_t(nsm)*2*MMQ_FEW_NROWS ? 2 : 1;
+}
+
+template <ggml_type type>
+static void launch_mul_mat_q_few_tall(const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst,
+        const int4 * src1_q, cudaStream_t stream) {
+    if constexpr (mmq_few_tall_supported(type)) {
+        if (mmq_few_nmt(type, 2, src0->ne[1]) == 2) {
+            launch_mul_mat_q_few<type, 2, 2>(src0, src1, dst, src1_q, stream);
+            return;
+        }
+    }
+    launch_mul_mat_q_few<type, 2, 1>(src0, src1, dst, src1_q, stream);
+}
+
+template <ggml_type type>
+static void launch_mul_mat_q_few_switch_ncg(const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst,
+        const int4 * src1_q, const int ncg, cudaStream_t stream) {
+    switch (ncg) {
+        case 1:
+            launch_mul_mat_q_few<type, 1, 1>(src0, src1, dst, src1_q, stream);
+            break;
+        case 2:
+            launch_mul_mat_q_few_tall<type>(src0, src1, dst, src1_q, stream);
+            break;
+        case 3:
+            launch_mul_mat_q_few<type, 3, 1>(src0, src1, dst, src1_q, stream);
+            break;
+        default:
+            GGML_ABORT("fatal error");
+    }
+}
+
+void ggml_cuda_mul_mat_q_few(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+    GGML_TENSOR_BINARY_OP_LOCALS;
+
+    GGML_ASSERT(src1->type == GGML_TYPE_F32);
+    GGML_ASSERT( dst->type == GGML_TYPE_F32);
+    GGML_ASSERT(nb00 == ggml_type_size(src0->type));
+    GGML_ASSERT(nb10 == sizeof(float));
+    GGML_ASSERT(nb0  == sizeof(float));
+    GGML_ASSERT(ne11 <= MMQ_FEW_MAX_NCG*MMQ_FEW_NCOLS);
+    GGML_ASSERT(ne00 % QK_K == 0);
+    GGML_ASSERT(ne12 % ne02 == 0 && ne13 % ne03 == 0);
+
+    cudaStream_t stream = ctx.stream();
+
+    const int nunits = ne00 / QK_K;
+    const int ncg    = mmq_few_ncg(ne11);
+
+    ggml_cuda_pool_alloc<int4> src1_q(ctx.pool(), int64_t(ne12)*ne13*nunits*ncg*MMQ_FEW_Q8_1_INT4);
+    quantize_mmq_few_q8_1_cuda((const float *) src1->data, src1_q.get(), mmq_few_ds_layout(src0->type), ne10, nb11/sizeof(float), nb12/sizeof(float), nb13/sizeof(float),
+        ne11, ne12, ne13, stream);
+
+    switch (src0->type) {
+        case GGML_TYPE_Q8_0:
+            launch_mul_mat_q_few_switch_ncg<GGML_TYPE_Q8_0>(src0, src1, dst, src1_q.get(), ncg, stream);
+            break;
+        case GGML_TYPE_Q3_K:
+            launch_mul_mat_q_few_switch_ncg<GGML_TYPE_Q3_K>(src0, src1, dst, src1_q.get(), ncg, stream);
+            break;
+        case GGML_TYPE_Q4_K:
+            launch_mul_mat_q_few_switch_ncg<GGML_TYPE_Q4_K>(src0, src1, dst, src1_q.get(), ncg, stream);
+            break;
+        case GGML_TYPE_Q5_K:
+            launch_mul_mat_q_few_switch_ncg<GGML_TYPE_Q5_K>(src0, src1, dst, src1_q.get(), ncg, stream);
+            break;
+        case GGML_TYPE_Q6_K:
+            launch_mul_mat_q_few_switch_ncg<GGML_TYPE_Q6_K>(src0, src1, dst, src1_q.get(), ncg, stream);
+            break;
+        case GGML_TYPE_IQ3_S:
+            launch_mul_mat_q_few_switch_ncg<GGML_TYPE_IQ3_S>(src0, src1, dst, src1_q.get(), ncg, stream);
+            break;
+        case GGML_TYPE_IQ4_NL:
+            launch_mul_mat_q_few_switch_ncg<GGML_TYPE_IQ4_NL>(src0, src1, dst, src1_q.get(), ncg, stream);
+            break;
+        case GGML_TYPE_IQ4_XS:
+            launch_mul_mat_q_few_switch_ncg<GGML_TYPE_IQ4_XS>(src0, src1, dst, src1_q.get(), ncg, stream);
+            break;
+        default:
+            GGML_ABORT("unsupported type for few-row MMQ");
+    }
+}
+
+// Min. number of src1 columns for the few-row MMQ, 0 if it is not used for this type, GPU and row length.
+static int ggml_cuda_mmq_few_min_ncols(const ggml_type type, const int cc, const int64_t ne00) {
+#ifdef GGML_CUDA_FORCE_CUBLAS
+    return 0;
+#endif // GGML_CUDA_FORCE_CUBLAS
+
+    // NVIDIA Ampere to Blackwell
+    if (!ampere_mma_available(cc) || cc >= GGML_CUDA_CC_RUBIN) {
+        return 0;
+    }
+    // on DGX Spark the kernel needs its L2 prefetch, which is only in the DGX Spark device code
+    if (cc == GGML_CUDA_CC_DGX_SPARK && ggml_cuda_highest_compiled_arch(cc) != GGML_CUDA_CC_DGX_SPARK) {
+        return 0;
+    }
+    // with fewer than 8 units of 256 values per row the warps idle, MMVQ and MMQ are better there
+    if (mmq_few_unit_size(type) == 0 || ne00 < 8*QK_K) {
+        return 0;
+    }
+    if (cc == GGML_CUDA_CC_DGX_SPARK) { // tuned on DGX Spark GB10
+        return type == GGML_TYPE_IQ4_XS || type == GGML_TYPE_IQ4_NL ? 8 : 2;
+    }
+    switch (type) { // all GPUs except DGX Spark, tuned on RTX 5090
+        case GGML_TYPE_Q3_K:   return 3;
+        case GGML_TYPE_Q5_K:   return 3;
+        case GGML_TYPE_Q4_K:   return 4;
+        case GGML_TYPE_IQ4_NL: return 4;
+        case GGML_TYPE_Q6_K:   return 4;
+        case GGML_TYPE_Q8_0:   return 5;
+        case GGML_TYPE_IQ3_S:  return 5;
+        case GGML_TYPE_IQ4_XS: return 6;
+        default:               return 0;
+    }
+}
+
+bool ggml_cuda_should_use_mmq_few(const ggml_tensor * src0, const ggml_tensor * src1, const int cc) {
+    const int min_ncols = ggml_cuda_mmq_few_min_ncols(src0->type, cc, src0->ne[0]);
+    if (min_ncols == 0 || src1->ne[1] < min_ncols || src1->ne[1] > MMQ_FEW_MAX_NCG*MMQ_FEW_NCOLS) {
+        return false;
+    }
+    // the row must hold a whole number of ring stages of the launched kernel
+    const int ncg          = mmq_few_ncg(src1->ne[1]);
+    const int nmt          = mmq_few_nmt(src0->type, ncg, src0->ne[1]);
+    const int nunits_stage = mmq_few_nunits_stage(src0->type, ncg, nmt);
+    if (src0->ne[0] % (nunits_stage*QK_K) != 0) {
+        return false;
+    }
+    // unaligned data needs the alignment of the decoder loads
+    const int x_alignment = mmq_few_unaligned(src0->type, nunits_stage) ? (ggml_type_size(src0->type) % 4 == 0 ? 4 : 2) : 16;
+    return ggml_cuda_is_aligned(src0, x_alignment) && ggml_cuda_is_aligned(src1, 16);
+}
+
+#else // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+
+void ggml_cuda_mul_mat_q_few(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+    GGML_UNUSED_VARS(ctx, src0, src1, dst);
+    GGML_ABORT("few-row MMQ is not supported on HIP or MUSA");
+}
+
+bool ggml_cuda_should_use_mmq_few(const ggml_tensor * src0, const ggml_tensor * src1, const int cc) {
+    GGML_UNUSED_VARS(src0, src1, cc);
+    return false;
+}
+
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)

ggml/src/ggml-cuda/mmq-few.cuh

@@ -0,0 +1,7 @@
+#pragma once
+
+#include "common.cuh"
+
+bool ggml_cuda_should_use_mmq_few(const ggml_tensor * src0, const ggml_tensor * src1, int cc);
+
+void ggml_cuda_mul_mat_q_few(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst);

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

@@ -629,6 +629,13 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     }
 }
 
+// The 6-bit scales 4*ksc .. 4*ksc + 3 of a q3_K block, minus 32, as int8.
+static __device__ __forceinline__ int unpack_scales_q3_K(const uint8_t * scales, const int ksc) {
+    const int sc_low  = (get_int_b2(scales, ksc % 2) >> (4 * (ksc / 2))) & 0x0F0F0F0F;
+    const int sc_high = ((get_int_b2(scales, 2) >> (2 * ksc)) << 4) & 0x30303030;
+    return __vsub4(sc_low | sc_high, 0x20202020);
+}
+
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q3_K(
         const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
     GGML_UNUSED(x_gate);
@@ -695,15 +702,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 
         const int ksc = threadIdx.x % 4;
 
-        const int ksc_low = ksc % (QI3_K/8);
-        const int shift_low = 4 * (ksc / (QI3_K/8));
-        const int sc_low = (get_int_b2(bxi->scales, ksc_low) >> shift_low) & 0x0F0F0F0F;
-
-        const int ksc_high = QI3_K/8;
-        const int shift_high = 2 * ksc;
-        const int sc_high = ((get_int_b2(bxi->scales, ksc_high) >> shift_high) << 4) & 0x30303030;
-
-        const int sc = __vsub4(sc_low | sc_high, 0x20202020);
+        const int sc = unpack_scales_q3_K(bxi->scales, ksc);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
         const int8_t * sc8 = (const int8_t *) &sc;

ggml/src/ggml-cuda/mmq.cu

@@ -169,21 +169,27 @@ void ggml_cuda_mul_mat_q(
     GGML_ASSERT(!ids || ids->type  == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID.
 
     // Fused gate/up/GLU: src0 is the up weight, dst the GLU output.
+    // Fused scale (no gate): dst is the MUL that applies the weight scale to the MUL_MAT.
     const ggml_tensor * gate = fusion ? fusion->gate : nullptr;
     ggml_cuda_mm_fusion_args_device fusion_device = {};
     if (fusion) {
-        GGML_ASSERT(gate && !fusion->x_bias && !fusion->gate_bias && !fusion->x_scale == !fusion->gate_scale);
-        GGML_ASSERT(gate->type == src0->type && ggml_are_same_shape(gate, src0) && ggml_are_same_stride(gate, src0));
+        GGML_ASSERT(!fusion->x_bias && !fusion->gate_bias);
         // NVFP4 weight scales (only applied for NVFP4): one value per tensor, or one per expert with ids.
         for (const ggml_tensor * scale : { fusion->x_scale, fusion->gate_scale }) {
             GGML_ASSERT(!scale || (src0->type == GGML_TYPE_NVFP4 && scale->type == GGML_TYPE_F32 && ggml_is_contiguous(scale) &&
                 ggml_nelements(scale) == (ids ? src0->ne[2] : 1)));
         }
-        fusion_device.gate       = gate->data;
-        fusion_device.x_scale    = fusion->x_scale    ? fusion->x_scale->data    : nullptr;
-        fusion_device.gate_scale = fusion->gate_scale ? fusion->gate_scale->data : nullptr;
-        fusion_device.glu_op     = fusion->glu_op;
-        fusion_device.glu_limit  = fusion->glu_limit;
+        if (gate) {
+            GGML_ASSERT(!fusion->x_scale == !fusion->gate_scale);
+            GGML_ASSERT(gate->type == src0->type && ggml_are_same_shape(gate, src0) && ggml_are_same_stride(gate, src0));
+            fusion_device.gate       = gate->data;
+            fusion_device.gate_scale = fusion->gate_scale ? fusion->gate_scale->data : nullptr;
+            fusion_device.glu_op     = fusion->glu_op;
+            fusion_device.glu_limit  = fusion->glu_limit;
+        } else {
+            GGML_ASSERT(!fusion->gate_scale && fusion->x_scale);
+        }
+        fusion_device.x_scale = fusion->x_scale ? fusion->x_scale->data : nullptr;
     }
 
     GGML_TENSOR_BINARY_OP_LOCALS;
@@ -231,8 +237,8 @@ void ggml_cuda_mul_mat_q(
 
     const bool fallback = ggml_cuda_mmq_needs_fallback(ne01);
 
-    // With fusion dst is the GLU output, not a MUL_MAT; the src1 precision requested by both fused MUL_MATs is in the fusion args.
-    const ggml_prec prec_src1 = ggml_cuda_mmq_get_prec_src1(src0->type, gate ? fusion->prec_src1 : ggml_prec_get_src(dst, 1), cc);
+    // With fusion dst is the GLU output or the scale node, not a MUL_MAT; the src1 precision requested by the fused MUL_MATs is in the fusion args.
+    const ggml_prec prec_src1 = ggml_cuda_mmq_get_prec_src1(src0->type, fusion ? fusion->prec_src1 : ggml_prec_get_src(dst, 1), cc);
 
     const bool use_native_fp4 = prec_src1 == GGML_PREC_Q4;
     const size_t y_block_size       = use_native_fp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq);

ggml/src/ggml-cuda/mmq.cuh

@@ -465,7 +465,11 @@ template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __de
         const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, float * __restrict__ dst,
         const float * __restrict__ y_scale, const int stride, const int i_max, const int j_sup,
         const ggml_glu_op glu_op, const float glu_limit, const float * __restrict__ x_scale, const float * __restrict__ gate_scale) {
-    GGML_UNUSED_VARS(glu_op, glu_limit, x_scale, gate_scale);
+    GGML_UNUSED_VARS(glu_op, glu_limit, gate_scale);
+
+    // NVFP4 weight scale of a fused MUL (per tensor, or of the expert of this tile), 1 without.
+    const float scale = type == GGML_TYPE_NVFP4 && x_scale ? *x_scale : 1.0f;
+
     constexpr int warp_size = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
     constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
@@ -489,11 +493,15 @@ template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __de
                 continue;
             }
 
-            if (y_scale_used) {
-                dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
+            float value = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
+            if constexpr (type == GGML_TYPE_NVFP4) {
+                if (y_scale_used) {
+                    value *= y_scale[j];
+                }
             } else {
-                dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
+                GGML_UNUSED(y_scale_used);
             }
+            dst[ids_dst[j]*stride + i] = value * scale;
         }
     }
 }
@@ -520,9 +528,9 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
 
     const bool y_scale_used = type == GGML_TYPE_NVFP4 && y_scale != nullptr;
 
-    // Fused NVFP4: load the weight scales once, 1 without scales.
+    // NVFP4: load the weight scales once, 1 without scales. Without fusion x_scale is the scale of a fused MUL.
     constexpr bool fused_nvfp4 = has_fusion && type == GGML_TYPE_NVFP4;
-    const float scale_up   = fused_nvfp4 && x_scale ? *x_scale    : 1.0f;
+    const float scale_up   = type == GGML_TYPE_NVFP4 && x_scale ? *x_scale : 1.0f;
     const float scale_gate = fused_nvfp4 && gate_scale ? *gate_scale : 1.0f;
 
 #pragma unroll
@@ -556,10 +564,16 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
                         gate *= scale_gate;
                     }
                     dst[ids_dst[j]*stride + i] = ggml_cuda_op_glu_single(glu_op, gate, up, glu_limit);
-                } else if (y_scale_used) {
-                    dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l];
                 } else {
-                    dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
+                    float value = sum[(j0/tile_C::J + n)*tile_C::ne + l];
+                    if constexpr (type == GGML_TYPE_NVFP4) {
+                        if (y_scale_used) {
+                            value *= y_scale[j];
+                        }
+                    } else {
+                        GGML_UNUSED(y_scale_used);
+                    }
+                    dst[ids_dst[j]*stride + i] = value * scale_up;
                 }
             }
         }
@@ -1201,10 +1215,12 @@ static __global__ void mul_mat_q(
 
         const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
 
+        const float * x_scale_tile = x_scale && ids_dst ? x_scale + zt : x_scale;
+
         constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer.
         mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1, false>
             (x, x_gate, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
-             glu_op, glu_limit, x_scale, gate_scale, stride_row_x, ncols_y, stride_col_dst,
+             glu_op, glu_limit, x_scale_tile, gate_scale, stride_row_x, ncols_y, stride_col_dst,
              tile_x_i_max, col_high - col_low, kb0_start, kb0_stop);
 
         kbc += blocks_per_ne00.z;
@@ -1275,10 +1291,12 @@ static __global__ void mul_mat_q(
 
     const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
 
+    const float * x_scale_tile = x_scale && ids_dst ? x_scale + zt : x_scale;
+
     constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks.
     mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1, false>
         (x, x_gate, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
-         glu_op, glu_limit, x_scale, gate_scale, stride_row_x, ncols_y, stride_col_dst,
+         glu_op, glu_limit, x_scale_tile, gate_scale, stride_row_x, ncols_y, stride_col_dst,
          tile_x_i_max, col_high - col_low, kb0_start, kb0_stop);
 }
 

ggml/src/ggml-cuda/mmvq.cu

@@ -29,10 +29,6 @@ static constexpr __host__ __device__ bool mmvq_should_prefetch(ggml_type type) {
             return false;
     }
 }
-
-static __device__ __forceinline__ void mmvq_prefetch_l2(const void * p) {
-    asm volatile("prefetch.global.L2 [%0];" :: "l"(p));
-}
 #endif
 
 typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs);
@@ -757,10 +753,10 @@ static __global__ void mul_mat_vec_q(
 #pragma unroll
                 for (int i = 0; i < rows_per_cuda_block; ++i) {
                     const size_t off = (size_t)(kbx_offset + i*stride_row_x + kbx_pf) * ggml_cuda_type_traits<type>::bs;
-                    mmvq_prefetch_l2((const char *) vx + off);
+                    ggml_cuda_prefetch_l2((const char *) vx + off);
                     if constexpr (has_fusion) {
                         if (use_gate) {
-                            mmvq_prefetch_l2((const char *) vgate + off);
+                            ggml_cuda_prefetch_l2((const char *) vgate + off);
                         }
                     }
                 }

ggml/src/ggml-cuda/quantize.cu

@@ -600,6 +600,81 @@ void quantize_mmq_q8_1_cuda(
     }
 }
 
+// Each thread quantizes 4 values of one column, see quantize.cuh for the layout.
+template <mmq_q8_1_ds_layout ds_layout>
+static __global__ void quantize_mmq_few_q8_1(
+        const float * x_ptr, void * vy_ptr, const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+        const int ne1, const int ne2) {
+    ggml_cuda_pdl_lc();
+    const float * GGML_CUDA_RESTRICT x  = x_ptr;
+    int4        * GGML_CUDA_RESTRICT vy = (int4 *) vy_ptr;
+    const int64_t i0 = ((int64_t)blockDim.x*blockIdx.y + threadIdx.x)*4;
+
+    if (i0 >= ne00) {
+        return;
+    }
+    ggml_cuda_pdl_sync();
+
+    const int64_t i2 = blockIdx.z % ne2;
+    const int64_t i3 = blockIdx.z / ne2;
+    const float4  xi = (int) blockIdx.x < ne1 ? *(const float4 *) (x + i3*s03 + i2*s02 + blockIdx.x*s01 + i0) : make_float4(0.0f, 0.0f, 0.0f, 0.0f);
+
+    // the 32 values of a sub-block are held by 8 threads
+    float amax = fmaxf(fmaxf(fabsf(xi.x), fabsf(xi.y)), fmaxf(fabsf(xi.z), fabsf(xi.w)));
+    amax = warp_reduce_max<8>(amax);
+
+    const float d_inv = 127.0f / amax;
+    char4 q;
+    q.x = roundf(xi.x*d_inv);
+    q.y = roundf(xi.y*d_inv);
+    q.z = roundf(xi.z*d_inv);
+    q.w = roundf(xi.w*d_inv);
+    const float d = 1.0f / d_inv;
+
+    // as in quantize_mmq_q8_1: the sum of the quantized values; the int warp_reduce_sum reduces the whole warp on Ampere+
+    [[maybe_unused]] float sum = 0.0f;
+    if constexpr (ds_layout != MMQ_Q8_1_DS_LAYOUT_D4) {
+        sum = d*warp_reduce_sum<8>(float(q.x + q.y + q.z + q.w));
+    }
+
+    const int s = i0 % QK_K / QK8_1;          // sub-block
+    const int k = i0 % QK8_1;                 // first value in sub-block
+    const int g = blockIdx.x % MMQ_FEW_NCOLS; // column in group
+    int4 * yg = vy + ((blockIdx.z*(ne00/QK_K) + i0/QK_K)*(gridDim.x/MMQ_FEW_NCOLS) + blockIdx.x/MMQ_FEW_NCOLS)*MMQ_FEW_Q8_1_INT4;
+    ((int *) (yg + (s/2)*WARP_SIZE + 4*g + (k % 16)/4))[2*(s % 2) + k/16] = *(const int *) &q;
+    if (k == 0) {
+        int4 * ds = yg + 4*WARP_SIZE + 4*(g/2) + s/2;
+        if constexpr (ds_layout == MMQ_Q8_1_DS_LAYOUT_D4) {
+            ((float *) ds)[2*(s % 2) + g % 2] = d;
+        } else {
+            ((half2 *) ds)[2*(s % 2) + g % 2] = make_half2(d, sum);
+        }
+    }
+}
+
+void quantize_mmq_few_q8_1_cuda(
+        const float * x, void * vy, const mmq_q8_1_ds_layout ds_layout, const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
+        const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
+    GGML_ASSERT(ne00 % QK_K == 0);
+
+    const int64_t block_num_y = (ne00 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ);
+    const dim3 num_blocks(GGML_PAD(ne1, MMQ_FEW_NCOLS), block_num_y, ne2*ne3);
+    const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
+    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(num_blocks, block_size, 0, stream);
+    switch (ds_layout) {
+        case MMQ_Q8_1_DS_LAYOUT_D4:
+            ggml_cuda_kernel_launch(quantize_mmq_few_q8_1<MMQ_Q8_1_DS_LAYOUT_D4>, launch_params,
+                x, vy, ne00, s01, s02, s03, (int) ne1, (int) ne2);
+            break;
+        case MMQ_Q8_1_DS_LAYOUT_DS4:
+            ggml_cuda_kernel_launch(quantize_mmq_few_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4>, launch_params,
+                x, vy, ne00, s01, s02, s03, (int) ne1, (int) ne2);
+            break;
+        default:
+            GGML_ABORT("fatal error");
+    }
+}
+
 // scatter=true reuses the quant kernel: grid over tokens, ids = inverse map (token slot -> compact row)
 void quantize_scatter_mmq_q8_1_cuda(
         const float * x, const int32_t * ids_src1_inv, void * vy, const ggml_type type_src0,

ggml/src/ggml-cuda/quantize.cuh

@@ -21,6 +21,17 @@ void quantize_mmq_q8_1_cuda(
         ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
         int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream);
 
+// src1 of the few-row MMQ in the layout of the MMA fragments. Per channel, 256 values and group of MMQ_FEW_NCOLS columns:
+//   qs: int4[4][32], [j][lane] holds the B fragments of sub-blocks 2j and 2j+1 of lane (g, t): the values [4t, 4t + 4) and [16 + 4t, 16 + 4t + 4) of column g.
+//   ds: int4[4][4], [t][j] holds d (D4, float) or (d, sum) (DS4, half2) of the columns 2t and 2t+1 for sub-block 2j, then 2j+1.
+// The column groups of a unit are adjacent, columns beyond ne1 are zero.
+static constexpr int MMQ_FEW_NCOLS     = 8; // src1 columns per column group
+static constexpr int MMQ_FEW_Q8_1_INT4 = 4*WARP_SIZE + 4*4; // int4 per 256 values and column group
+
+void quantize_mmq_few_q8_1_cuda(
+        const float * x, void * vy, mmq_q8_1_ds_layout ds_layout, int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
+        int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream);
+
 void quantize_mmq_fp4_cuda(const float *   x,
                              const int32_t * ids,
                              void *          vy,

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

@@ -3428,6 +3428,10 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
                   "please update hexagon_type to match ggml_type");
     static_assert((unsigned int) HTP_TYPE_Q4_1 == (unsigned int) GGML_TYPE_Q4_1,
                   "please update hexagon_type to match ggml_type");
+    static_assert((unsigned int) HTP_TYPE_Q5_0 == (unsigned int) GGML_TYPE_Q5_0,
+                  "please update hexagon_type to match ggml_type");
+    static_assert((unsigned int) HTP_TYPE_Q5_1 == (unsigned int) GGML_TYPE_Q5_1,
+                  "please update hexagon_type to match ggml_type");
     static_assert((unsigned int) HTP_TYPE_Q8_0 == (unsigned int) GGML_TYPE_Q8_0,
                   "please update hexagon_type to match ggml_type");
     static_assert((unsigned int) HTP_TYPE_MXFP4 == (unsigned int) GGML_TYPE_MXFP4,

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

@@ -953,6 +953,8 @@ static bool htp_matmul_check_types(
             case GGML_TYPE_Q2_0:
             case GGML_TYPE_Q4_0:
             case GGML_TYPE_Q4_1:
+            case GGML_TYPE_Q5_0:
+            case GGML_TYPE_Q5_1:
             case GGML_TYPE_Q8_0:
             case GGML_TYPE_IQ4_NL:
             case GGML_TYPE_MXFP4:
@@ -995,6 +997,8 @@ static bool htp_matmul_check_types(
         case GGML_TYPE_Q2_0:
         case GGML_TYPE_Q4_0:
         case GGML_TYPE_Q4_1:
+        case GGML_TYPE_Q5_0:
+        case GGML_TYPE_Q5_1:
         case GGML_TYPE_Q8_0:
         case GGML_TYPE_IQ4_NL:
         case GGML_TYPE_MXFP4:

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

@@ -15,7 +15,7 @@
 #    include <arm_neon.h>
 #endif
 
-// --- 4-bit block-32 nibble helpers (shared by Q4_0, IQ4_NL, Q4_1, MXFP4) ---
+// --- 4-bit block-32 nibble helpers (shared by Q4_0, IQ4_NL, Q4_1, Q5_0, Q5_1, MXFP4) ---
 
 static inline void unpack_nibbles_32(uint8_t * dst, const uint8_t * qs) {
     for (unsigned int i = 0; i < 16; ++i) {
@@ -48,6 +48,28 @@ static inline void transpose_tile_to_nibbles(uint8_t tile_quants[32][32], const
     }
 }
 
+// 5th-bit plane of one row (qh bit k = quant k), same bit layout as the Q5_K plane
+static inline void pack_q5_plane(uint8_t * plane, int row, uint32_t qh) {
+    for (int c = 0; c < 4; c++) {
+        uint32_t v = 0;
+        for (int i = 0; i < 4; i++) {
+            v |= ((qh >> (8 * i + 2 * c)) & 3) << (2 * i);
+        }
+        plane[c * 32 + row] = (uint8_t) v;
+    }
+}
+
+static inline uint32_t unpack_q5_plane(const uint8_t * plane, int row) {
+    uint32_t qh = 0;
+    for (int c = 0; c < 4; c++) {
+        const uint32_t v = plane[c * 32 + row];
+        for (int i = 0; i < 4; i++) {
+            qh |= ((v >> (2 * i)) & 3) << (8 * i + 2 * c);
+        }
+    }
+    return qh;
+}
+
 #if defined(__ARM_NEON) || defined(__aarch64__)
 static void repack_neon_transpose8x8(const uint8x8_t rows[8], uint8x8_t columns[8]) {
     const uint8x8x2_t t0 = vtrn_u8(rows[0], rows[1]);
@@ -433,7 +455,7 @@ static void unpack_from_tiled_grid_32(void * data, const ggml_tensor * t, size_t
     }
 }
 
-// --- Block-32 Quantized Types (Q4_0, IQ4_NL, Q4_1, Q8_0, MXFP4) ---
+// --- Block-32 Quantized Types (Q4_0, IQ4_NL, Q4_1, Q5_0, Q5_1, Q8_0, MXFP4) ---
 
 static void repack_q4_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
     repack_to_tiled_grid_32<block_q4_0, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q4_0,
@@ -546,6 +568,141 @@ static void repack_tiled_q4_1(void * data, const ggml_tensor * t, size_t offset,
         });
 }
 
+static void repack_q5_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
+    repack_to_tiled_grid_32<block_q5_0, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q5_0,
+        [](uint8_t * tile_dst, const block_q5_0 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
+            uint8_t *   plane_dst = tile_dst + 512;
+            ggml_half * scale_dst = (ggml_half *) (tile_dst + 640);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+                for (int row0 = 0; row0 < 32; row0 += 8) {
+                    const uint8_t * qs[8];
+                    for (int rr = 0; rr < 8; ++rr) {
+                        qs[rr] = src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+                    }
+                    repack_neon_4bit_rows(qs, tile_dst, row0);
+                }
+
+                for (int row = 0; row < 32; ++row) {
+                    const block_q5_0 & src_block = src_slice[(ct * 32 + row) * row_stride + kt];
+                    uint32_t qh;
+                    memcpy(&qh, src_block.qh, sizeof(qh));
+                    pack_q5_plane(plane_dst, row, ~qh);
+                    scale_dst[row] = src_block.d;
+                }
+                return;
+            }
+#endif
+            uint8_t tile_quants[32][32];
+            for (int row = 0; row < 32; row++) {
+                const int64_t r = ct * 32 + row;
+                if (r < ne1 && kt < ne0 / 32) {
+                    const block_q5_0 & b = src_slice[r * row_stride + kt];
+                    unpack_nibbles_32(tile_quants[row], b.qs);
+                    uint32_t qh;
+                    memcpy(&qh, b.qh, sizeof(qh));
+                    pack_q5_plane(plane_dst, row, ~qh);
+                    scale_dst[row] = b.d;
+                } else {
+                    memset(tile_quants[row], 0, 32);
+                    pack_q5_plane(plane_dst, row, 0);
+                    scale_dst[row] = 0;
+                }
+            }
+            transpose_nibbles_to_tile(tile_dst, tile_quants);
+        });
+}
+
+static void repack_tiled_q5_0(void * data, const ggml_tensor * t, size_t offset, size_t size) {
+    unpack_from_tiled_grid_32<block_q5_0, 1>(data, t, offset, size, HTP_TENSOR_TILE_SIZE_Q5_0,
+        [](block_q5_0 * dst_slice, const uint8_t * tile_src, int64_t ne0, int64_t start_row, int64_t end_row, int64_t row_stride, int ct, int kt) {
+            uint8_t tile_quants[32][32];
+            transpose_tile_to_nibbles(tile_quants, tile_src);
+            const uint8_t *   plane_src = tile_src + 512;
+            const ggml_half * scale_src = (const ggml_half *) (tile_src + 640);
+            for (int row = 0; row < 32; row++) {
+                const int64_t r = ct * 32 + row;
+                if (r >= start_row && r < end_row && kt < ne0 / 32) {
+                    block_q5_0 & b = dst_slice[(r - start_row) * row_stride + kt];
+                    pack_nibbles_32(b.qs, tile_quants[row]);
+                    const uint32_t qh = ~unpack_q5_plane(plane_src, row);
+                    memcpy(b.qh, &qh, sizeof(qh));
+                    b.d = scale_src[row];
+                }
+            }
+        });
+}
+
+// Q5_1 uses the Q5_K tile layout
+static void repack_q5_1_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
+    repack_to_tiled_grid_32<block_q5_1, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q5_1,
+        [](uint8_t * tile_dst, const block_q5_1 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
+            ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
+            uint8_t *   plane_dst = tile_dst + 640;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+                for (int row0 = 0; row0 < 32; row0 += 8) {
+                    const uint8_t * qs[8];
+                    for (int rr = 0; rr < 8; ++rr) {
+                        qs[rr] = src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+                    }
+                    repack_neon_4bit_rows(qs, tile_dst, row0);
+                }
+
+                for (int row = 0; row < 32; ++row) {
+                    const block_q5_1 & src_block = src_slice[(ct * 32 + row) * row_stride + kt];
+                    uint32_t qh;
+                    memcpy(&qh, src_block.qh, sizeof(qh));
+                    pack_q5_plane(plane_dst, row, qh);
+                    scale_dst[2 * row + 0] = src_block.d;
+                    scale_dst[2 * row + 1] = src_block.m;
+                }
+                return;
+            }
+#endif
+            uint8_t tile_quants[32][32];
+            for (int row = 0; row < 32; row++) {
+                const int64_t r = ct * 32 + row;
+                if (r < ne1 && kt < ne0 / 32) {
+                    const block_q5_1 & b = src_slice[r * row_stride + kt];
+                    unpack_nibbles_32(tile_quants[row], b.qs);
+                    uint32_t qh;
+                    memcpy(&qh, b.qh, sizeof(qh));
+                    pack_q5_plane(plane_dst, row, qh);
+                    scale_dst[2 * row + 0] = b.d;
+                    scale_dst[2 * row + 1] = b.m;
+                } else {
+                    memset(tile_quants[row], 0, 32);
+                    pack_q5_plane(plane_dst, row, 0);
+                    scale_dst[2 * row + 0] = 0;
+                    scale_dst[2 * row + 1] = 0;
+                }
+            }
+            transpose_nibbles_to_tile(tile_dst, tile_quants);
+        });
+}
+
+static void repack_tiled_q5_1(void * data, const ggml_tensor * t, size_t offset, size_t size) {
+    unpack_from_tiled_grid_32<block_q5_1, 1>(data, t, offset, size, HTP_TENSOR_TILE_SIZE_Q5_1,
+        [](block_q5_1 * dst_slice, const uint8_t * tile_src, int64_t ne0, int64_t start_row, int64_t end_row, int64_t row_stride, int ct, int kt) {
+            uint8_t tile_quants[32][32];
+            transpose_tile_to_nibbles(tile_quants, tile_src);
+            const ggml_half * scale_src = (const ggml_half *) (tile_src + 512);
+            const uint8_t *   plane_src = tile_src + 640;
+            for (int row = 0; row < 32; row++) {
+                const int64_t r = ct * 32 + row;
+                if (r >= start_row && r < end_row && kt < ne0 / 32) {
+                    block_q5_1 & b = dst_slice[(r - start_row) * row_stride + kt];
+                    pack_nibbles_32(b.qs, tile_quants[row]);
+                    const uint32_t qh = unpack_q5_plane(plane_src, row);
+                    memcpy(b.qh, &qh, sizeof(qh));
+                    b.d = scale_src[2 * row + 0];
+                    b.m = scale_src[2 * row + 1];
+                }
+            }
+        });
+}
+
 static void repack_mxfp4_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
     repack_to_tiled_grid_32<block_mxfp4, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_MXFP4,
         [](uint8_t * tile_dst, const block_mxfp4 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
@@ -1538,6 +1695,12 @@ void htp_repack_tensor_tiled(ggml_tensor * tensor, const void * data, size_t siz
         case GGML_TYPE_Q4_1:
             repack_q4_1_tiled(tensor, data, 0, size);
             break;
+        case GGML_TYPE_Q5_0:
+            repack_q5_0_tiled(tensor, data, 0, size);
+            break;
+        case GGML_TYPE_Q5_1:
+            repack_q5_1_tiled(tensor, data, 0, size);
+            break;
         case GGML_TYPE_Q8_0:
             repack_q8_0_tiled(tensor, data, 0, size);
             break;
@@ -1593,6 +1756,12 @@ void htp_unpack_tensor_tiled(void * data, const ggml_tensor * tensor, size_t off
         case GGML_TYPE_Q4_1:
             repack_tiled_q4_1(data, tensor, offset, size);
             break;
+        case GGML_TYPE_Q5_0:
+            repack_tiled_q5_0(data, tensor, offset, size);
+            break;
+        case GGML_TYPE_Q5_1:
+            repack_tiled_q5_1(data, tensor, offset, size);
+            break;
         case GGML_TYPE_Q8_0:
             repack_tiled_q8_0(data, tensor, offset, size);
             break;

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

@@ -10,6 +10,7 @@
 static inline bool htp_is_tiled_weight_type(enum ggml_type type) {
     return type == GGML_TYPE_Q1_0  || type == GGML_TYPE_Q2_0   ||
            type == GGML_TYPE_Q4_0  || type == GGML_TYPE_Q4_1   ||
+           type == GGML_TYPE_Q5_0  || type == GGML_TYPE_Q5_1   ||
            type == GGML_TYPE_MXFP4 || type == GGML_TYPE_IQ4_NL ||
            type == GGML_TYPE_Q8_0  || type == GGML_TYPE_Q6_K   ||
            type == GGML_TYPE_Q2_K  || type == GGML_TYPE_Q3_K   ||

ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h

@@ -647,6 +647,106 @@ static void dequantize_tiled_weight_to_fp16_task_q5_k(
     }
 }
 
+static void dequantize_tiled_weight_to_fp16_task_q5_0(
+        const tiled_dequantize_state_t *state,
+        uint32_t start_tile, uint32_t end_tile) {
+
+    const HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F);
+
+    for (uint32_t t = start_tile; t < end_tile; t++) {
+        const uint8_t * tile_src = state->src + t * state->aligned_tile_size;
+        __fp16 * dst_ptr = state->dst + t * HTP_MM_HMX_TILE_N_ELMS;
+
+        HVX_Vector v_sc = hvx_vmem(tile_src + 640);
+        HVX_Vector v_scale_duplicated = Q6_V_lo_W(Q6_W_vshuff_VVR(v_sc, v_sc, -2));
+
+        // Load all 4 groups in parallel
+        HVX_Vector vq0 = hvx_vmem(tile_src + 0 * 128);
+        HVX_Vector vq1 = hvx_vmem(tile_src + 1 * 128);
+        HVX_Vector vq2 = hvx_vmem(tile_src + 2 * 128);
+        HVX_Vector vq3 = hvx_vmem(tile_src + 3 * 128);
+
+        // Nibble extraction
+        HVX_Vector v_lo0 = Q6_V_vand_VV(vq0, mask_h4);
+        HVX_Vector v_hi0 = Q6_Vub_vlsr_VubR(vq0, 4);
+        HVX_Vector v_lo1 = Q6_V_vand_VV(vq1, mask_h4);
+        HVX_Vector v_hi1 = Q6_Vub_vlsr_VubR(vq1, 4);
+        HVX_Vector v_lo2 = Q6_V_vand_VV(vq2, mask_h4);
+        HVX_Vector v_hi2 = Q6_Vub_vlsr_VubR(vq2, 4);
+        HVX_Vector v_lo3 = Q6_V_vand_VV(vq3, mask_h4);
+        HVX_Vector v_hi3 = Q6_Vub_vlsr_VubR(vq3, 4);
+
+        // Q5_0: OR 0xF0 where the plane flags a clear 5th bit (q - 16)
+        HVX_Vector v_plane = hvx_vmem(tile_src + 512);
+        v_lo0 = hvx_q5_0_or_neg(v_lo0, v_plane, 0);
+        v_hi0 = hvx_q5_0_or_neg(v_hi0, v_plane, 1);
+        v_lo1 = hvx_q5_0_or_neg(v_lo1, v_plane, 2);
+        v_hi1 = hvx_q5_0_or_neg(v_hi1, v_plane, 3);
+        v_lo2 = hvx_q5_0_or_neg(v_lo2, v_plane, 4);
+        v_hi2 = hvx_q5_0_or_neg(v_hi2, v_plane, 5);
+        v_lo3 = hvx_q5_0_or_neg(v_lo3, v_plane, 6);
+        v_hi3 = hvx_q5_0_or_neg(v_hi3, v_plane, 7);
+
+        // Shuffling
+        HVX_VectorPair vp_shuf0 = Q6_W_vshuff_VVR(v_hi0, v_lo0, -1);
+        HVX_VectorPair vp_shuf1 = Q6_W_vshuff_VVR(v_hi1, v_lo1, -1);
+        HVX_VectorPair vp_shuf2 = Q6_W_vshuff_VVR(v_hi2, v_lo2, -1);
+        HVX_VectorPair vp_shuf3 = Q6_W_vshuff_VVR(v_hi3, v_lo3, -1);
+
+        // Unpack to 16-bit
+        HVX_VectorPair vp_int16_lo0 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf0));
+        HVX_VectorPair vp_int16_hi0 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf0));
+        HVX_VectorPair vp_int16_lo1 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf1));
+        HVX_VectorPair vp_int16_hi1 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf1));
+        HVX_VectorPair vp_int16_lo2 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf2));
+        HVX_VectorPair vp_int16_hi2 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf2));
+        HVX_VectorPair vp_int16_lo3 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf3));
+        HVX_VectorPair vp_int16_hi3 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf3));
+
+        // Convert and scale multiplication
+        HVX_Vector v_grp0_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo0)), v_scale_duplicated));
+        HVX_Vector v_grp0_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo0)), v_scale_duplicated));
+        HVX_Vector v_grp0_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi0)), v_scale_duplicated));
+        HVX_Vector v_grp0_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi0)), v_scale_duplicated));
+
+        HVX_Vector v_grp1_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo1)), v_scale_duplicated));
+        HVX_Vector v_grp1_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo1)), v_scale_duplicated));
+        HVX_Vector v_grp1_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi1)), v_scale_duplicated));
+        HVX_Vector v_grp1_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi1)), v_scale_duplicated));
+
+        HVX_Vector v_grp2_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo2)), v_scale_duplicated));
+        HVX_Vector v_grp2_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo2)), v_scale_duplicated));
+        HVX_Vector v_grp2_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi2)), v_scale_duplicated));
+        HVX_Vector v_grp2_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi2)), v_scale_duplicated));
+
+        HVX_Vector v_grp3_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo3)), v_scale_duplicated));
+        HVX_Vector v_grp3_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo3)), v_scale_duplicated));
+        HVX_Vector v_grp3_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi3)), v_scale_duplicated));
+        HVX_Vector v_grp3_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi3)), v_scale_duplicated));
+
+        // Parallel Stores
+        hvx_vmem(dst_ptr +  0 * 64) = v_grp0_0;
+        hvx_vmem(dst_ptr +  1 * 64) = v_grp0_1;
+        hvx_vmem(dst_ptr +  2 * 64) = v_grp0_2;
+        hvx_vmem(dst_ptr +  3 * 64) = v_grp0_3;
+
+        hvx_vmem(dst_ptr +  4 * 64) = v_grp1_0;
+        hvx_vmem(dst_ptr +  5 * 64) = v_grp1_1;
+        hvx_vmem(dst_ptr +  6 * 64) = v_grp1_2;
+        hvx_vmem(dst_ptr +  7 * 64) = v_grp1_3;
+
+        hvx_vmem(dst_ptr +  8 * 64) = v_grp2_0;
+        hvx_vmem(dst_ptr +  9 * 64) = v_grp2_1;
+        hvx_vmem(dst_ptr + 10 * 64) = v_grp2_2;
+        hvx_vmem(dst_ptr + 11 * 64) = v_grp2_3;
+
+        hvx_vmem(dst_ptr + 12 * 64) = v_grp3_0;
+        hvx_vmem(dst_ptr + 13 * 64) = v_grp3_1;
+        hvx_vmem(dst_ptr + 14 * 64) = v_grp3_2;
+        hvx_vmem(dst_ptr + 15 * 64) = v_grp3_3;
+    }
+}
+
 // Q6_K stores 6-bit weights and one fp16 scale per 16 k, see HTP_TENSOR_TILE_SIZE_Q6_K.
 // A k-group holds 4 k per row, the HMX tile holds 2, so each group is dealt into two tiles.
 static void dequantize_tiled_weight_to_fp16_task_q6_k(

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

@@ -21,6 +21,8 @@ enum htp_data_type {
     HTP_TYPE_F16    = 1,
     HTP_TYPE_Q4_0   = 2,
     HTP_TYPE_Q4_1   = 3,
+    HTP_TYPE_Q5_0   = 6,
+    HTP_TYPE_Q5_1   = 7,
     HTP_TYPE_Q8_0   = 8,
     HTP_TYPE_Q2_K   = 10,
     HTP_TYPE_Q3_K   = 11,

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

@@ -17,6 +17,8 @@
 #define HTP_TENSOR_TILE_SIZE_F16    2048
 #define HTP_TENSOR_TILE_SIZE_Q4_0   576
 #define HTP_TENSOR_TILE_SIZE_Q4_1   640
+#define HTP_TENSOR_TILE_SIZE_Q5_0   704
+#define HTP_TENSOR_TILE_SIZE_Q5_1   768
 #define HTP_TENSOR_TILE_SIZE_Q8_0   1088
 #define HTP_TENSOR_TILE_SIZE_IQ4_NL 576
 #define HTP_TENSOR_TILE_SIZE_MXFP4  544
@@ -30,7 +32,7 @@
 #define HTP_TENSOR_TILE_K_Q2_0      64
 #define HTP_TENSOR_TILE_SIZE_Q2_0   576
 
-// Tiled weight tensor layouts (Q4_0, Q8_0, IQ4_NL, MXFP4, Q1_0, Q2_0, K-quants):
+// Tiled weight tensor layouts (Q4_0, Q5_0, Q5_1, Q8_0, IQ4_NL, MXFP4, Q1_0, Q2_0, K-quants):
 // See docs/backend/snapdragon/developer.md ("Tiled Weight Formats").
 
 struct htp_tensor_mdev_range {
@@ -143,6 +145,8 @@ static inline uint32_t htp_tensor_get_tile_size(int type) {
         case HTP_TYPE_Q4_0:   return HTP_TENSOR_TILE_SIZE_Q4_0;
         case HTP_TYPE_IQ4_NL: return HTP_TENSOR_TILE_SIZE_IQ4_NL;
         case HTP_TYPE_Q4_1:   return HTP_TENSOR_TILE_SIZE_Q4_1;
+        case HTP_TYPE_Q5_0:   return HTP_TENSOR_TILE_SIZE_Q5_0;
+        case HTP_TYPE_Q5_1:   return HTP_TENSOR_TILE_SIZE_Q5_1;
         case HTP_TYPE_Q4_K:   return HTP_TENSOR_TILE_SIZE_Q4_K;
         case HTP_TYPE_Q8_0:   return HTP_TENSOR_TILE_SIZE_Q8_0;
         case HTP_TYPE_Q5_K:   return HTP_TENSOR_TILE_SIZE_Q5_K;
@@ -182,6 +186,8 @@ static inline uint32_t htp_tensor_get_row_size(int type, uint32_t ne00) {
         case HTP_TYPE_Q4_0:
         case HTP_TYPE_IQ4_NL: return (ne00 / 32) * 18;
         case HTP_TYPE_Q4_1:   return (ne00 / 32) * 20;
+        case HTP_TYPE_Q5_0:   return (ne00 / 32) * 22;
+        case HTP_TYPE_Q5_1:   return (ne00 / 32) * 24;
         case HTP_TYPE_MXFP4:  return (ne00 / 32) * 17;
         case HTP_TYPE_Q2_K:   return (ne00 / 256) * 84;
         case HTP_TYPE_Q2_0:   return (ne00 / 64) * 18;

ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h

@@ -639,6 +639,74 @@ static inline HVX_VectorPair accum_5bit_32x2(
     return Q6_W_vcombine_VV(v_sum1, v_sum0);
 }
 
+// OR 0xF0 into every lane of v whose flag j is set in the plane, i.e. 4-bit value | 0xF0 = q - 16 where the 5th bit is clear (Q5_0)
+static inline HVX_Vector hvx_q5_0_or_neg(HVX_Vector v, HVX_Vector v_plane, int j) {
+    HVX_VectorPred q = Q6_Q_vand_VR(v_plane, 0x01010101u << j);
+    return Q6_V_vandor_VQR(v, q, 0xF0F0F0F0);
+}
+
+// Q5_0: signed bytes (q - 16) straight from the nibbles and the plane
+static inline HVX_VectorPair unpack_and_interleave_q5_0_x2(HVX_Vector v_src, HVX_Vector v_plane, int i, HVX_Vector mask_h4) {
+    HVX_Vector v_lo = hvx_q5_0_or_neg(Q6_V_vand_VV(v_src, mask_h4), v_plane, 2 * i);
+    HVX_Vector v_hi = hvx_q5_0_or_neg(Q6_Vub_vlsr_VubR(v_src, 4), v_plane, 2 * i + 1);
+    HVX_VectorPair v01_pair = Q6_W_vshuff_VVR(v_hi, v_lo, -1);
+    HVX_Vector v01_lo = Q6_V_lo_W(v01_pair);
+    HVX_Vector v01_hi = Q6_V_hi_W(v01_pair);
+
+    HVX_Vector v23_lo = Q6_V_valign_VVR(v01_hi, v01_lo, 64);
+    HVX_Vector v_W0 = Q6_V_lo_W(Q6_W_vshuff_VVR(v23_lo, v01_lo, -2));
+
+    HVX_Vector v67_lo = Q6_V_valign_VVR(v01_lo, v01_hi, 64);
+    HVX_Vector v_W1 = Q6_V_lo_W(Q6_W_vshuff_VVR(v67_lo, v01_hi, -2));
+
+    return Q6_W_vcombine_VV(v_W1, v_W0);
+}
+
+static inline HVX_Vector accum_q5_0_32x1(
+    const HVX_Vector * restrict vptr,
+    const HVX_Vector * restrict v_act
+) {
+    HVX_Vector v_sum0 = Q6_V_vzero();
+    HVX_Vector v_sum1 = Q6_V_vzero();
+    HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F);
+    HVX_Vector v_plane = vptr[4];
+
+    #pragma unroll
+    for (int i = 0; i < 4; i++) {
+        HVX_VectorPair v_W_pair = unpack_and_interleave_q5_0_x2(vptr[i], v_plane, i, mask_h4);
+        v_sum0 = Q6_Vw_vrmpyacc_VwVbVb(v_sum0, Q6_V_lo_W(v_W_pair), v_act[i * 2 + 0]);
+        v_sum1 = Q6_Vw_vrmpyacc_VwVbVb(v_sum1, Q6_V_hi_W(v_W_pair), v_act[i * 2 + 1]);
+    }
+
+    return Q6_Vw_vadd_VwVw(v_sum0, v_sum1);
+}
+
+static inline HVX_VectorPair accum_q5_0_32x2(
+    const HVX_Vector * restrict vptr,
+    const HVX_Vector * restrict v_act0,
+    const HVX_Vector * restrict v_act1
+) {
+    HVX_Vector v_sum0 = Q6_V_vzero();
+    HVX_Vector v_sum1 = Q6_V_vzero();
+    HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F);
+    HVX_Vector v_plane = vptr[4];
+
+    #pragma unroll
+    for (int i = 0; i < 4; i++) {
+        HVX_VectorPair v_W_pair = unpack_and_interleave_q5_0_x2(vptr[i], v_plane, i, mask_h4);
+        HVX_Vector v_W0 = Q6_V_lo_W(v_W_pair);
+        HVX_Vector v_W1 = Q6_V_hi_W(v_W_pair);
+
+        v_sum0 = Q6_Vw_vrmpyacc_VwVbVb(v_sum0, v_W0, v_act0[i * 2 + 0]);
+        v_sum0 = Q6_Vw_vrmpyacc_VwVbVb(v_sum0, v_W1, v_act0[i * 2 + 1]);
+
+        v_sum1 = Q6_Vw_vrmpyacc_VwVbVb(v_sum1, v_W0, v_act1[i * 2 + 0]);
+        v_sum1 = Q6_Vw_vrmpyacc_VwVbVb(v_sum1, v_W1, v_act1[i * 2 + 1]);
+    }
+
+    return Q6_W_vcombine_VV(v_sum1, v_sum0);
+}
+
 // Q6_K weights are stored unsigned (0..63), see HTP_TENSOR_TILE_SIZE_Q6_K. Unpack k-group g of a tile to signed bytes (q - 32)
 static inline HVX_Vector unpack_q6_k_group(const HVX_Vector * restrict vptr, int g, HVX_Vector mask_0f, HVX_Vector mask_03, HVX_Vector i32) {
     HVX_Vector v_lo = (g & 1) ? Q6_Vub_vlsr_VubR(vptr[g >> 1], 4) : Q6_V_vand_VV(vptr[g >> 1], mask_0f);
@@ -1251,6 +1319,82 @@ static void tiled_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float
     }
 }
 
+static void tiled_vec_dot_q5_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
+    const uint8_t * restrict tile_ptr = vx;
+    const uint8_t * restrict y_q = vy;
+
+    HVX_Vector v_sum_float = Q6_V_vzero();
+
+    uint32_t n_k_tiles = n / 32;
+    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
+        const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 768);
+        const HVX_Vector * restrict v_act = (const HVX_Vector *) (y_q + kt * 1152);
+
+        HVX_Vector v_sum = accum_q5_0_32x1(vptr, v_act);
+        HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum);
+
+        HVX_Vector v_scale_w = vptr[5];
+        HVX_Vector v_scale_a = v_act[8];
+        HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a);
+        HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb);
+
+        v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
+    }
+
+    if (sz) {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+    } else {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    }
+}
+
+static void tiled_vec_dot_q5_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
+    const uint8_t * restrict tile_ptr = vx;
+    const uint8_t * restrict y0_q = vy0;
+    const uint8_t * restrict y1_q = vy1;
+
+    HVX_Vector v_sum_float_c0 = Q6_V_vzero();
+    HVX_Vector v_sum_float_c1 = Q6_V_vzero();
+
+    uint32_t n_k_tiles = n / 32;
+    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
+        const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 768);
+        const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1152);
+        const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1152);
+
+        HVX_VectorPair v_sums = accum_q5_0_32x2(vptr, v_act0, v_act1);
+        HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums);
+        HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums);
+
+        HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0);
+        HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1);
+
+        HVX_Vector v_scale_w = vptr[5];
+        HVX_Vector v_scale_a_c0 = v_act0[8];
+        HVX_Vector v_scale_a_c1 = v_act1[8];
+
+        HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a_c0);
+        HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a_c1);
+
+        HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0);
+        HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1);
+
+        v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0);
+        v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
+    }
+
+    if (sz0) {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+    } else {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+    }
+    if (sz1) {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+    } else {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    }
+}
+
 static void tiled_vec_dot_q5_k_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y_q = vy;

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

@@ -1664,6 +1664,12 @@ static int hvx_mm_init_tiled_vec_dot(struct htp_mm_context * mmctx, enum htp_dat
             mmctx->vec_dot_32x1 = tiled_vec_dot_q8_0_32x1;
             mmctx->vec_dot_32x2 = tiled_vec_dot_q8_0_32x2;
             return 0;
+        case HTP_TYPE_Q5_0:
+            mmctx->type         = "q5_0_tiled-f32";
+            mmctx->vec_dot_32x1 = tiled_vec_dot_q5_0_32x1;
+            mmctx->vec_dot_32x2 = tiled_vec_dot_q5_0_32x2;
+            return 0;
+        case HTP_TYPE_Q5_1:
         case HTP_TYPE_Q5_K:
             mmctx->type         = "q5_k_tiled-f32";
             mmctx->vec_dot_32x1 = tiled_vec_dot_q5_k_32x1;
@@ -2316,6 +2322,8 @@ static inline tiled_dequant_task_fn_t hmx_get_dequant_worker_fn(int weight_type,
         case HTP_TYPE_Q4_K:   return dequantize_tiled_weight_to_fp16_task_q4_1;
         case HTP_TYPE_MXFP4:  return dequantize_tiled_weight_to_fp16_task_mxfp4;
         case HTP_TYPE_Q8_0:   return dequantize_tiled_weight_to_fp16_task_q8_0;
+        case HTP_TYPE_Q5_0:   return dequantize_tiled_weight_to_fp16_task_q5_0;
+        case HTP_TYPE_Q5_1:
         case HTP_TYPE_Q5_K:   return dequantize_tiled_weight_to_fp16_task_q5_k;
         case HTP_TYPE_Q6_K:   return dequantize_tiled_weight_to_fp16_task_q6_k;
         case HTP_TYPE_Q3_K:   return dequantize_tiled_weight_to_fp16_task_q3_k;

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

@@ -185,7 +185,7 @@ next_nc:
 // (Q2_K: per-16 k sums, q8_1_s16)
 static inline bool htp_mm_weight_has_offset(int weight_type) {
     return weight_type == HTP_TYPE_Q4_1 || weight_type == HTP_TYPE_Q4_K || weight_type == HTP_TYPE_Q5_K ||
-           weight_type == HTP_TYPE_Q2_K || weight_type == HTP_TYPE_Q2_0;
+           weight_type == HTP_TYPE_Q2_K || weight_type == HTP_TYPE_Q2_0 || weight_type == HTP_TYPE_Q5_1;
 }
 
 // --- Activation/Row Size Helpers ---

ggml/src/ggml-opencl/CMakeLists.txt

@@ -76,7 +76,6 @@ set(GGML_OPENCL_KERNELS
     clamp
     cpy
     cvt
-    diag_mask_inf
     diag
     div
     dsv4_hc

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

@@ -798,7 +798,6 @@ struct ggml_backend_opencl_context {
     cl_program program_moe_add_id_glu;
     cl_program program_clamp;
     cl_program program_cvt;
-    cl_program program_diag_mask_inf;
     cl_program program_gelu;
     cl_program program_gemv_noshuffle_general;
     cl_program program_gemv_noshuffle;
@@ -886,7 +885,6 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_rms_norm_mul_add = nullptr;   // fused rms_norm(x)*w + b (residual)
     cl_kernel kernel_l2_norm_f32;
     cl_kernel kernel_group_norm, kernel_group_norm_mul_add;
-    cl_kernel kernel_diag_mask_inf, kernel_diag_mask_inf_8;
     cl_kernel kernel_diag_f32;
     cl_kernel kernel_soft_max, kernel_soft_max_4;
     cl_kernel kernel_soft_max_f16, kernel_soft_max_4_f16;
@@ -2291,23 +2289,6 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
         GGML_LOG_CONT(".");
     }
 
-    // diag_mask_inf
-    {
-#ifdef GGML_OPENCL_EMBED_KERNELS
-        const std::string kernel_src {
-            #include "diag_mask_inf.cl.h"
-        };
-#else
-        const std::string kernel_src = read_file("diag_mask_inf.cl");
-#endif
-        backend_ctx->program_diag_mask_inf =
-            build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
-
-        CL_CHECK((backend_ctx->kernel_diag_mask_inf_8 = clCreateKernel(backend_ctx->program_diag_mask_inf, "kernel_diag_mask_inf_8", &err), err));
-        CL_CHECK((backend_ctx->kernel_diag_mask_inf   = clCreateKernel(backend_ctx->program_diag_mask_inf, "kernel_diag_mask_inf", &err), err));
-        GGML_LOG_CONT(".");
-    }
-
     // diag
     {
 #ifdef GGML_OPENCL_EMBED_KERNELS
@@ -12214,8 +12195,6 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
             return true;
         case GGML_OP_DIAG:
             return true;
-        case GGML_OP_DIAG_MASK_INF:
-            return op->ne[3] == 1;
         case GGML_OP_ROPE: {
             const int mode = ((const int32_t *) op->op_params)[2];
             if (!ggml_rope_get_truncate(op)) {
@@ -37948,68 +37927,6 @@ static void ggml_cl_set(ggml_backend_t backend, const ggml_tensor * src0, const
     backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
 }
 
-static void ggml_cl_diag_mask_inf(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
-    GGML_ASSERT(src0);
-    GGML_ASSERT(src0->extra);
-    GGML_ASSERT(dst);
-    GGML_ASSERT(dst->extra);
-
-    UNUSED(src1);
-
-    int n_past = ((int32_t *)(dst->op_params))[0];
-
-    const int  ne00 = src0 ? src0->ne[0] : 0;
-    const int  ne01 = src0 ? src0->ne[1] : 0;
-    const int  ne02 = src0 ? src0->ne[2] : 0;
-
-    ggml_backend_opencl_context *backend_ctx = (ggml_backend_opencl_context *)backend->context;
-
-    ggml_tensor_extra_cl * extra0 = (ggml_tensor_extra_cl *)src0->extra;
-    ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *)dst->extra;
-
-    cl_ulong offset0 = extra0->offset + src0->view_offs;
-    cl_ulong offsetd = extrad->offset + dst->view_offs;
-
-    cl_kernel kernel;
-
-    if (ne00%8 == 0) {
-        kernel = backend_ctx->kernel_diag_mask_inf_8;
-
-        CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem),   &extra0->data_device));
-        CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset0));
-        CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem),   &extrad->data_device));
-        CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offsetd));
-        CL_CHECK(clSetKernelArg(kernel, 4, sizeof(int),      &ne00));
-        CL_CHECK(clSetKernelArg(kernel, 5, sizeof(int),      &ne01));
-        CL_CHECK(clSetKernelArg(kernel, 6, sizeof(int),      &n_past));
-
-        size_t global_work_size[] = {(size_t)ne00*ne01*ne02/8, 1, 1};
-        size_t local_work_size[] = {64, 1, 1};
-
-        backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
-    } else {
-        kernel = backend_ctx->kernel_diag_mask_inf;
-
-        CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem),   &extra0->data_device));
-        CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset0));
-        CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem),   &extrad->data_device));
-        CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offsetd));
-        CL_CHECK(clSetKernelArg(kernel, 4, sizeof(int),      &ne00));
-        CL_CHECK(clSetKernelArg(kernel, 5, sizeof(int),      &ne01));
-        CL_CHECK(clSetKernelArg(kernel, 6, sizeof(int),      &n_past));
-
-        size_t global_work_size[] = {(size_t)ne00, (size_t)ne01, (size_t)ne02};
-        size_t local_work_size[] = {64, 1, 1};
-
-        size_t * local_work_size_ptr = local_work_size;
-        if (ne00 % 64 != 0 && !backend_ctx->non_uniform_workgroups) {
-            local_work_size_ptr = nullptr;  // Let driver choose the work-group sizes.
-        }
-
-        backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size_ptr, dst);
-    }
-}
-
 static void ggml_cl_diag(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
     GGML_ASSERT(src0);
     GGML_ASSERT(src0->extra);
@@ -39935,12 +39852,6 @@ bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor
             }
             func = ggml_cl_diag;
             break;
-        case GGML_OP_DIAG_MASK_INF:
-            if (!any_on_device) {
-                return false;
-            }
-            func = ggml_cl_diag_mask_inf;
-            break;
         case GGML_OP_SOFT_MAX:
             if (!any_on_device) {
                 return false;

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

@@ -1,58 +0,0 @@
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-
-//------------------------------------------------------------------------------
-// diag_mask_inf kernels
-//------------------------------------------------------------------------------
-kernel void kernel_diag_mask_inf(
-        global float * src0,
-        ulong offset0,
-        global float * dst,
-        ulong offsetd,
-        int ne00,
-        int ne01,
-        int n_past
-) {
-    src0 = (global float*)((global char*)src0 + offset0);
-    dst = (global float*)((global char*)dst + offsetd);
-
-    int i02 = get_global_id(2);
-    int i01 = get_global_id(1);
-    int i00 = get_global_id(0);
-
-    if (i00 > n_past + i01) {
-        dst[i02*ne01*ne00 + i01*ne00 + i00] = -INFINITY;
-    } else {
-        dst[i02*ne01*ne00 + i01*ne00 + i00] = src0[i02*ne01*ne00 + i01*ne00 + i00];
-    }
-}
-
-kernel void kernel_diag_mask_inf_8(
-        global float4 * src0,
-        ulong offset0,
-        global float4 * dst,
-        ulong offsetd,
-        int ne00,
-        int ne01,
-        int n_past
-) {
-    src0 = (global float4*)((global char*)src0 + offset0);
-    dst = (global float4*)((global char*)dst + offsetd);
-
-    int i = 2*get_global_id(0);
-
-    dst[i+0] = src0[i+0];
-    dst[i+1] = src0[i+1];
-    int i4 = 4*i;
-    int i02 = i4/(ne00*ne01); i4 -= i02*ne00*ne01;
-    int i01 = i4/(ne00);      i4 -= i01*ne00;
-    int i00 = i4;
-    for (int k = 3; k >= 0; --k) {
-        if (i00 + 4 + k <= n_past + i01) {
-            break;
-        }
-        (&dst[i+1])[k] = -INFINITY;
-        if (i00 + k > n_past + i01) {
-            (&dst[i])[k] = -INFINITY;
-        }
-    }
-}

ggml/src/ggml-openvino/openvino/op_table.cpp

@@ -33,7 +33,6 @@ std::unordered_map<std::string, CreatorFunction> get_supported_ops() {
     using namespace ov::op;
     return {
         {"GGML_OP_ADD",             op::translate_add                              },
-        {"GGML_OP_ADD1",            op::translate_1to1_match_2_inputs<v1::Add>     },
         {"GGML_OP_ADD_ID",          op::translate_add_id                           },
         {"GGML_OP_CONCAT",          op::translate_concat                           },
         {"GGML_OP_CONT",            op::translate_cont                             },

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

@@ -2517,23 +2517,6 @@ k_argsort_f32_i32(const float *x, int *dst, const int ncols, int ncols_pad,
     }
 }
 
-static void diag_mask_inf_f32(const float * x, float * dst, const int ncols, const int rows_per_channel, const int n_past,
-                              const sycl::nd_item<3> &item_ct1) {
-    const int col = item_ct1.get_local_range(1) * item_ct1.get_group(1) +
-                    item_ct1.get_local_id(1);
-    const int row = item_ct1.get_local_range(2) * item_ct1.get_group(2) +
-                    item_ct1.get_local_id(2);
-
-    if (col >= ncols) {
-        return;
-    }
-
-    const int i = row*ncols + col;
-    //dst[i] = col > (n_past + row % rows_per_channel) ? -INFINITY : x[i];
-    //dst[i] = x[i] - (col > n_past + row % rows_per_channel) * INT_MAX; // equivalent within rounding error but slightly faster on GPU
-    dst[i] = x[i] - (col > n_past + row % rows_per_channel) * FLT_MAX;
-}
-
 static void scale_f32(const float * x, float * dst, const float scale, const float bias, const int k,
                       const sycl::nd_item<3> &item_ct1) {
     const int i = item_ct1.get_local_range(2) * item_ct1.get_group(2) +
@@ -3139,20 +3122,6 @@ static void argmax_f32_i32_sycl(const float *x, int *dst, const int ncols,
             });
     });
 }
-static void diag_mask_inf_f32_sycl(const float *x, float *dst,
-                                   const int ncols_x, const int nrows_x,
-                                   const int rows_per_channel, const int n_past,
-                                   queue_ptr stream) {
-    const sycl::range<3> block_dims(1, SYCL_DIAG_MASK_INF_BLOCK_SIZE, 1);
-    const int block_num_x = (ncols_x + SYCL_DIAG_MASK_INF_BLOCK_SIZE - 1) / SYCL_DIAG_MASK_INF_BLOCK_SIZE;
-    const sycl::range<3> block_nums(1, block_num_x, nrows_x);
-    stream->parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
-                         [=](sycl::nd_item<3> item_ct1) {
-                             diag_mask_inf_f32(x, dst, ncols_x,
-                                               rows_per_channel, n_past,
-                                               item_ct1);
-                         });
-}
 
 static dpct::err0 ggml_sycl_cpy_tensor_2d(void *dst,
                                           const struct ggml_tensor *src,
@@ -3536,25 +3505,6 @@ inline void ggml_sycl_op_argmax(ggml_backend_sycl_context & ctx, ggml_tensor * d
     argmax_f32_i32_sycl(src0_dd, dst_dd, ncols, nrows, main_stream);
 }
 
-inline void ggml_sycl_op_diag_mask_inf(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
-    GGML_ASSERT(dst->src[0]->type == GGML_TYPE_F32);
-    GGML_ASSERT( dst->type == GGML_TYPE_F32);
-    dpct::queue_ptr main_stream = ctx.stream();
-    SYCL_CHECK(ggml_sycl_set_device(ctx.device));
-    const float * src0_dd = static_cast<const float *>(dst->src[0]->data);
-    float *       dst_dd  = static_cast<float *>(dst->data);
-
-    const int64_t ne00 = dst->src[0]->ne[0];
-    const int64_t ne01 = dst->src[0]->ne[1];
-    const int nrows0 = ggml_nrows(dst->src[0]);
-    // the kernel indexes in 32-bit
-    GGML_ASSERT(ne00 * nrows0 <= INT32_MAX);
-
-    const int n_past = ((int32_t *) dst->op_params)[0];
-
-    diag_mask_inf_f32_sycl(src0_dd, dst_dd, ne00, nrows0, ne01, n_past, main_stream);
-}
-
 static void tri_f32_sycl(
     const float * src,
     float * dst,
@@ -6045,11 +5995,6 @@ static void ggml_sycl_scale(ggml_backend_sycl_context & ctx, ggml_tensor * dst)
     ggml_sycl_op_scale(ctx, dst);
 }
 
-static void ggml_sycl_diag_mask_inf(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
-    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/1);
-    ggml_sycl_op_diag_mask_inf(ctx, dst);
-}
-
 static void ggml_sycl_pool2d(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/1);
     ggml_sycl_op_pool2d(ctx, dst);
@@ -6193,7 +6138,6 @@ static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct gg
             ggml_sycl_dup(ctx, dst);
             break;
         case GGML_OP_ADD:
-        case GGML_OP_ADD1: // TODO: more efficient implementation
             ggml_sycl_add(ctx, dst);
             break;
         case GGML_OP_ADD_ID:
@@ -6396,9 +6340,6 @@ static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct gg
         case GGML_OP_TRI:
             ggml_sycl_op_tri(ctx, dst);
             break;
-        case GGML_OP_DIAG_MASK_INF:
-            ggml_sycl_diag_mask_inf(ctx, dst);
-            break;
         case GGML_OP_SOFT_MAX:
             ggml_sycl_op_soft_max(ctx, dst);
             break;
@@ -7639,7 +7580,6 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
         case GGML_OP_PERMUTE:
         case GGML_OP_TRANSPOSE:
         case GGML_OP_ADD:
-        case GGML_OP_ADD1:
         case GGML_OP_ADD_ID:
         case GGML_OP_SUB:
         case GGML_OP_COUNT_EQUAL:
@@ -7674,8 +7614,6 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
                        op->type == GGML_TYPE_F32 &&
                        ggml_is_contiguous(src0);
             }
-        case GGML_OP_DIAG_MASK_INF:
-            return true;
         case GGML_OP_SOFT_MAX:
             return true;
         case GGML_OP_SOFT_MAX_BACK: {

ggml/src/ggml-sycl/presets.hpp

@@ -35,7 +35,6 @@
 #define SYCL_SCALE_BLOCK_SIZE 256
 #define SYCL_ROPE_BLOCK_SIZE 256
 #define SYCL_ALIBI_BLOCK_SIZE 32
-#define SYCL_DIAG_MASK_INF_BLOCK_SIZE 32
 #define SYCL_QUANTIZE_BLOCK_SIZE 256
 #define SYCL_DEQUANTIZE_BLOCK_SIZE 256
 #define SYCL_GET_ROWS_BLOCK_SIZE 256

ggml/src/ggml-vulkan/ggml-vulkan-common.h

@@ -172,7 +172,6 @@ void ggml_vk_l2_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const gg
 void ggml_vk_unary(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst);
 void ggml_vk_xielu(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst);
 void ggml_vk_glu(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst);
-void ggml_vk_diag_mask_inf(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst);
 void ggml_vk_soft_max(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst);
 void ggml_vk_soft_max_back(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst);
 void ggml_vk_topk_moe(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_cgraph * cgraph, int node_idx);

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

@@ -1035,8 +1035,6 @@ void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph
         } else if (tensor->op == GGML_OP_SCALE) {
             const float * params = (const float *)tensor->op_params;
             tensor_clone = ggml_scale_bias(ggml_ctx, src_clone[0], params[0], params[1]);
-        } else if (tensor->op == GGML_OP_ADD1) {
-            tensor_clone = ggml_add1(ggml_ctx, src_clone[0], src_clone[1]);
         } else if (tensor->op == GGML_OP_ARANGE) {
             const float start = ggml_get_op_params_f32(tensor, 0);
             const float stop = ggml_get_op_params_f32(tensor, 1);
@@ -1101,8 +1099,6 @@ void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph
             }
         } else if (tensor->op == GGML_OP_SOFT_MAX_BACK) {
             tensor_clone = ggml_soft_max_ext_back(ggml_ctx, src_clone[0], src_clone[1], ((float *)tensor->op_params)[0], ((float *)tensor->op_params)[1]);
-        } else if (tensor->op == GGML_OP_DIAG_MASK_INF) {
-            tensor_clone = ggml_diag_mask_inf(ggml_ctx, src_clone[0], tensor->op_params[0]);
         } else if (tensor->op == GGML_OP_ROPE || tensor->op == GGML_OP_ROPE_BACK) {
             const int n_dims      = ((int32_t *) tensor->op_params)[1];
             const int mode        = ((int32_t *) tensor->op_params)[2];

ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h

@@ -124,22 +124,30 @@ struct vk_flash_attn_push_constants {
 
 static_assert(sizeof(vk_flash_attn_push_constants) <= 128, "sizeof(vk_flash_attn_push_constants) must be <= 128");
 
-struct vk_fa_xe_opt_push_constants {
-    uint32_t kv_seq_len;
-    uint32_t activation_length;
-    uint32_t q_head;
-    uint32_t kv_head;
-    uint32_t qk_ratio;
-    uint32_t qk_sub_groups;
+struct vk_xe_fa_push_constants {
+    uint32_t neq1;
+    uint32_t neq2;
+    uint32_t neq3;
+    uint32_t nek1;
+    uint32_t nek2;
+    uint32_t nek3;
+    uint32_t nev2;
+    uint32_t nev3;
+    uint32_t nem1;
+    uint32_t nem2;
+    uint32_t nem3;
+    uint32_t nbq1;
+    uint32_t nbq2;
+    uint32_t nbk1;
+    uint32_t nbk2;
+    uint32_t nbv1;
+    uint32_t nbv2;
+    uint32_t nbq3;
+    uint32_t nbk3;
+    uint32_t nbv3;
+    uint32_t nb3;
     uint32_t flag;
-    uint32_t nbkv_tok;
-    uint32_t nbkv_head;
-    uint32_t batch_stride_q;
-    uint32_t batch_stride_k;
-    uint32_t batch_stride_v;
-    uint32_t batch_stride_m;
-    uint32_t batch_stride_o;
-    float softmax_scale;
+    float scale;
 };
 
 struct vk_op_push_constants {
@@ -454,12 +462,6 @@ struct vk_op_add_id_push_constants {
     uint32_t s21;
 };
 
-struct vk_op_diag_mask_push_constants {
-    uint32_t ncols;
-    uint32_t rows_per_channel;
-    int32_t n_past;
-};
-
 struct vk_op_rope_push_constants {
     uint32_t rope_mode;
     uint32_t nrows;

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

@@ -433,6 +433,21 @@ struct vk_fa_pipeline_state {
     }
 };
 
+struct vk_xe_fa_pipeline_state {
+    uint32_t HSK, HSV;
+    uint32_t qk_ratio, pv_ratio;
+    uint32_t decode_n;
+    bool unaligned_seq_len;
+    bool unaligned_q;
+    bool partial_row_mask;
+    bool multi_head_mask;
+
+    bool operator<(const vk_xe_fa_pipeline_state &b) const {
+        return std::tie(HSK, HSV, qk_ratio, pv_ratio, decode_n, unaligned_seq_len, unaligned_q, partial_row_mask, multi_head_mask) <
+               std::tie(b.HSK, b.HSV, b.qk_ratio, b.pv_ratio, b.decode_n, b.unaligned_seq_len, b.unaligned_q, b.partial_row_mask, b.multi_head_mask);
+    }
+};
+
 struct vk_conv2d_pipeline_state {
     vk_conv2d_pipeline_state(uint32_t s0, uint32_t s1, uint32_t p0, uint32_t p1, uint32_t d0, uint32_t d1, uint32_t KW, uint32_t KH, uint32_t aligned)
         : s0(s0), s1(s1), p0(p0), p1(p1), d0(d0), d1(d1), KW(KW), KH(KH), aligned(aligned) {}
@@ -763,6 +778,9 @@ struct vk_device_struct {
     bool dot2_f16 {};
     bool ocp_fp4 {};
 
+    bool workgroup_memory_explicit_layout {};
+    bool workgroup_memory_explicit_layout_16bit {};
+
     bool pipeline_executable_properties_support {};
 
     bool device_fault {};
@@ -925,7 +943,6 @@ struct vk_device_struct {
 
     vk_pipeline pipeline_leaky_relu[2];
     vk_pipeline pipeline_silu_back_f32;
-    vk_pipeline pipeline_diag_mask_inf_f32;
     vk_pipeline pipeline_soft_max_f32, pipeline_soft_max_f32_f16;
     vk_pipeline pipeline_soft_max_f32_wg512, pipeline_soft_max_f32_f16_wg512;
     vk_pipeline pipeline_soft_max_back_f32;
@@ -1006,7 +1023,8 @@ struct vk_device_struct {
     bool fa_sparse_compact_use_subgroups;
 
     vk_pipeline pipeline_flash_attn_split_k_reduce;
-    std::map<std::tuple<uint32_t, uint32_t, uint32_t, uint32_t>, std::pair<vk_pipeline, vk_pipeline>> pipeline_xe_fa_decode_dual_phases;
+    std::map<vk_xe_fa_pipeline_state, std::pair<vk_pipeline, vk_pipeline>> pipeline_xe_fa_decode_dual_phases;
+    std::map<vk_xe_fa_pipeline_state, vk_pipeline> pipeline_xe_fa_prefill;
     vk_pipeline pipeline_count_experts;
 
     // [2] is for whether to take n_experts from spec constant (0) or push constant (1)
@@ -1028,6 +1046,7 @@ struct vk_device_struct {
     bool disable_host_visible_vidmem;
     bool allow_sysmem_fallback;
     bool disable_graph_optimize;
+    bool xe_fa_support {};
 
     std::unique_ptr<vk_memory_logger> memory_logger;
 
@@ -1407,6 +1426,34 @@ struct vk_fa_tuning_params {
     }
 };
 
+struct vk_xe_fa_tuning_params {
+    size_t shmem_sz;
+    std::array<uint32_t, 2> workgroup_size;
+    std::array<uint32_t, 2> subgroup_size;
+    uint32_t tm;
+    uint32_t tn;
+    uint32_t tk;
+
+    struct {
+        uint32_t wg_row;
+        uint32_t gqa_fold;
+        uint32_t q_in_reg;
+        uint32_t p_per_iter;
+        uint32_t sg_col;
+        uint32_t acc_count;
+        uint32_t pipeline_flag;
+    } prefill_params;
+
+    struct {
+        uint32_t p_chunk_ph1;
+        uint32_t out_per_wg_ph1;
+        uint32_t split_p_per_iter_ph2;
+        uint32_t out_row_per_wg_ph2;
+        uint32_t out_col_per_wg_ph2;
+        uint32_t pipeline_flag_ph2;
+    } decode_params;
+};
+
 struct GpuPipelineConfig {
     // GPU architecture identifier.
     // Example: vk_device_architecture::AMD_GCN

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

@@ -1434,6 +1434,186 @@ static std::vector<uint32_t> get_fa_spec_constants(const vk_fa_pipeline_state& s
     };
 }
 
+static vk_xe_fa_tuning_params get_xe_fa_tuning_params(const vk_device_struct * device, const vk_xe_fa_pipeline_state& pipeline_state, bool use_prefill) {
+    vk_xe_fa_tuning_params params {};
+    size_t shmem_sz = 0;
+    uint32_t sg_sz = 32u;
+    std::array<uint32_t, 3> matrix_mnk = { 16u, 16u, 16u };
+    bool intel_gpu = device->vendor_id == VK_VENDOR_ID_INTEL;
+
+    if (device->architecture == INTEL_XE1) {
+        sg_sz = 8u;
+        matrix_mnk = { 8u, 8u, 16u };
+    } else if (device->architecture == INTEL_XE2) {
+        sg_sz = 16u;
+        matrix_mnk = { 8u, 16u, 16u };
+    }
+
+    uint32_t tm = matrix_mnk[0];
+    uint32_t tn = matrix_mnk[1];
+    uint32_t tk = matrix_mnk[2];
+    uint32_t hsk = pipeline_state.HSK;
+    uint32_t hsv = pipeline_state.HSV;
+    uint32_t qk_ratio = pipeline_state.qk_ratio;
+    uint32_t pv_ratio = pipeline_state.pv_ratio;
+    uint32_t sg_per_gw;
+
+    params.tm = tm;
+    params.tn = tn;
+    params.tk = tk;
+    GGML_ASSERT(hsv % tm == 0);
+
+    if (use_prefill) {
+        sg_per_gw = sg_sz == 8 ? 32 : 16;
+        uint32_t wg_sz = sg_per_gw * sg_sz;
+        uint32_t output_div = hsk > 256 ? 8 : hsk > 160 ? 4 : hsk > 64 ? 2 : 1;
+        uint32_t p_per_iter = intel_gpu ? (hsk > 256 ? 16 : 32) : 16;
+        uint32_t q_in_reg = hsk > 512 ? hsk - 384 : std::min(128u, hsk);
+        uint32_t q_in_shmem = hsk - q_in_reg;
+        uint32_t wg_row = 256u / output_div;
+        uint32_t gqa_fold = (pipeline_state.multi_head_mask || qk_ratio != pv_ratio) ? 1 : 16;
+        uint32_t max_fold_rows = sg_sz <= 16 ? 16 : 8;
+
+        while (gqa_fold > 1 && (qk_ratio % gqa_fold != 0 || gqa_fold * output_div > max_fold_rows)) {
+            gqa_fold = gqa_fold / 2;
+        }
+
+        uint32_t mask_row = wg_row / gqa_fold;
+        uint32_t k_col = (hsk % 64 == 0 && hsk <= 512 && wg_sz <= 256) ? 8 : 4;
+        uint32_t k_loops = CEIL_DIV(hsk * p_per_iter, k_col * wg_sz);
+
+        bool k_in_shmem = hsk <= 576u;
+        bool v_in_shmem = hsk <= 128u && hsv <= 128u;
+        bool emf_in_shmem = wg_row / matrix_mnk[1] != sg_per_gw;
+        bool mask_in_shmem = mask_row <= 64;
+        size_t shmem_q_sz = q_in_shmem * wg_row * sizeof(ggml_fp16_t);
+        size_t shmem_oob_staging_sz = sg_per_gw * tm * tk * sizeof(ggml_fp16_t);
+        size_t shmem_mask_sz = mask_in_shmem ? (mask_row * p_per_iter * sizeof(ggml_fp16_t)) : 0;
+        size_t shmem_k_sz = k_in_shmem ? (k_loops * k_col * wg_sz * sizeof(ggml_fp16_t)) : shmem_oob_staging_sz;
+        size_t shmem_v_sz = v_in_shmem ? (hsv * p_per_iter * sizeof(ggml_fp16_t)) : 0;
+        size_t shmem_p_sz = wg_row * p_per_iter * sizeof(float);
+        size_t shmem_emf_sz = emf_in_shmem ? (wg_row * sizeof(float)) : 0;
+        size_t shmem_pv_sz = intel_gpu ? 0 : (sg_per_gw * tm * tn * sizeof(float));
+
+        // Reuse section sits at tail of shmem, reused with k + mask/kv oob staging/pv staging
+        size_t shmem_reuse_sz = shmem_k_sz + shmem_mask_sz;
+        // When v is not in shmem, its staging reuses the k + mask region; size the reuse section to the larger of the two
+        shmem_reuse_sz = v_in_shmem ? shmem_reuse_sz : std::max(shmem_reuse_sz, shmem_oob_staging_sz);
+        shmem_reuse_sz = std::max(shmem_pv_sz, shmem_reuse_sz);
+
+        // Q shmem size check for prologue staging is safe to bypass:
+        // 1. Prologue q staging shmem is released right after prologue, so it never overlaps other shmem sections
+        // 2. Extreme large hsk one-time q staging is capped within shader (max_q_reg_staging = shmem_size / 2 / Wr)
+        shmem_sz = shmem_q_sz + shmem_reuse_sz + shmem_v_sz + shmem_p_sz + shmem_emf_sz;
+
+        uint32_t sg_col = hsv / tm;
+        uint32_t pipeline_flag = 0;
+        sg_col = (sg_col + output_div - 1) / output_div;
+        sg_col = sg_col * tm;
+        if (pipeline_state.unaligned_seq_len) {
+            pipeline_flag |= 1 << 0;
+        }
+
+        if (pipeline_state.unaligned_q) {
+            pipeline_flag |= 1 << 1;
+        }
+
+        if (pipeline_state.partial_row_mask) {
+            pipeline_flag |= 1 << 2;
+        }
+
+        if (intel_gpu) {
+            pipeline_flag |= 1 << 3;
+        }
+
+        params.workgroup_size = { wg_sz };
+        params.subgroup_size = { sg_sz };
+        params.shmem_sz = shmem_sz;
+        params.prefill_params.pipeline_flag = pipeline_flag;
+        params.prefill_params.wg_row = wg_row;
+        params.prefill_params.gqa_fold = gqa_fold;
+        params.prefill_params.q_in_reg = q_in_reg;
+        params.prefill_params.p_per_iter = p_per_iter;
+        params.prefill_params.sg_col = sg_col;
+        params.prefill_params.acc_count = tm == 8 ? (sg_col % 16 == 0 ? 2 : 1) : 1;
+    } else {
+        auto upper_power_of_2 = [&](uint32_t in) {
+            GGML_ASSERT(in != 0);
+            if (in <= 1) return 1u;
+            uint32_t ret = in - 1;
+            ret |= ret >> 1;
+            ret |= ret >> 2;
+            ret |= ret >> 4;
+            ret |= ret >> 8;
+            ret |= ret >> 16;
+            return ret + 1;
+        };
+
+        size_t shmem_ph1_sz = 0;
+        size_t shmem_ph2_sz = 0;
+
+        // Phase 1 shmem size calculation
+        const uint32_t p_chunk_ph1 = 32;
+        // Phase 1 k shmem staging for extreme large hsk is capped within shader, bypass that check
+        uint32_t pmat_row = CEIL_DIV(qk_ratio, tn);
+        uint32_t aligned_q_len = upper_power_of_2(pipeline_state.decode_n);
+        uint32_t group_sz_ph1 = std::min(std::max(aligned_q_len * sg_sz, 64u), 256u);
+        uint32_t out_per_wg_ph1 = std::min(pipeline_state.decode_n, 256u / sg_sz);
+        size_t shmem_k_sz = group_sz_ph1 / sg_sz * tm * tk * sizeof(ggml_fp16_t);
+        size_t shmem_mask_sz = out_per_wg_ph1 * p_chunk_ph1 * sizeof(float);
+        size_t shmem_p_sz = p_chunk_ph1 * pmat_row * tn * sizeof(float);
+        // K shmem is reused with mask/p, calculate the max size
+        shmem_ph1_sz = std::max(shmem_k_sz, shmem_mask_sz + shmem_p_sz);
+
+        // Phase 2 shmem size calculation
+        uint32_t aligned_pv_ratio = upper_power_of_2(pv_ratio);
+        uint32_t split_p_per_iter_ph2 = (pv_ratio <= 32) ? 256 : 128;
+        uint32_t out_col_per_wg_ph2 = intel_gpu ? (pv_ratio > 16 ? 8 : 16) : 16;
+        uint32_t split_p_per_warp = 16;
+        uint32_t group_sz_ph2 = (split_p_per_iter_ph2 / split_p_per_warp) * sg_sz;
+        sg_per_gw = group_sz_ph2 / sg_sz;
+        uint32_t out_row_per_wg_ph2 = std::min(std::max(16u / aligned_pv_ratio, 1u), pipeline_state.decode_n);
+        uint32_t o_row = pv_ratio * out_row_per_wg_ph2;
+        uint32_t omat_row = CEIL_DIV(o_row, tn) * tn;
+        uint32_t global_max_loops = CEIL_DIV(o_row, sg_per_gw);
+        size_t shmem_global_max_sz = CEIL_DIV(global_max_loops * sg_per_gw, sg_sz) * sg_sz * sizeof(float);
+
+        uint32_t size_max = o_row * split_p_per_iter_ph2 / p_chunk_ph1;
+        size_t shmem_group_max_sz = CEIL_DIV(size_max, group_sz_ph2) * group_sz_ph2 * sizeof(float);
+
+        uint32_t p_loop_col = group_sz_ph2 * 4;
+        uint32_t p_loops = CEIL_DIV(o_row * split_p_per_iter_ph2, p_loop_col);
+        shmem_p_sz = std::max(p_loop_col * p_loops * sizeof(ggml_fp16_t), omat_row * split_p_per_iter_ph2 * sizeof(ggml_fp16_t));
+
+        size_t shmem_lf_sz = CEIL_DIV(sg_per_gw * p_loops, sg_sz) * sg_sz * sizeof(float);
+        size_t shmem_oob_staging_sz = sg_per_gw * tm * tk * sizeof(ggml_fp16_t);
+        size_t shmem_o_sz = sg_per_gw * out_col_per_wg_ph2 * omat_row * sizeof(float);
+        size_t shmem_reuse_sz = 2 * shmem_group_max_sz + std::max(shmem_oob_staging_sz, 2 * shmem_p_sz);
+
+        shmem_ph2_sz = shmem_lf_sz + shmem_global_max_sz;
+        // Shmem of p/group max is reused with output accumulation, calculate the max size
+        shmem_ph2_sz += std::max(shmem_reuse_sz, shmem_o_sz);
+        shmem_sz = std::max(shmem_ph1_sz, shmem_ph2_sz);
+
+        params.workgroup_size = { group_sz_ph1, group_sz_ph2 };
+        params.subgroup_size = { sg_sz, sg_sz };
+        params.shmem_sz = shmem_sz;
+        params.decode_params.p_chunk_ph1 = p_chunk_ph1;
+        params.decode_params.out_per_wg_ph1 = out_per_wg_ph1;
+        params.decode_params.split_p_per_iter_ph2 = split_p_per_iter_ph2;
+        params.decode_params.out_row_per_wg_ph2 = out_row_per_wg_ph2;
+        params.decode_params.out_col_per_wg_ph2 = out_col_per_wg_ph2;
+        params.decode_params.pipeline_flag_ph2 = pipeline_state.unaligned_seq_len ? (1u << 0) : 0;
+    }
+
+    if (shmem_sz > device->properties.limits.maxComputeSharedMemorySize) {
+        VK_LOG_DEBUG("get_xe_fa_tuning_params(HSK=" << hsk << ", HSV=" << hsv << ", qk_ratio=" << qk_ratio << ", pv_ratio=" << pv_ratio << "), "
+                     "shmem_sz = " << shmem_sz << ", exceeding max shmem size = " << device->properties.limits.maxComputeSharedMemorySize);
+    }
+
+    return params;
+}
+
 static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vector<uint32_t>& warptile, bool mul_mat_id, ggml_type src0_type) {
 
     uint32_t lut_size = 0;
@@ -3202,43 +3382,36 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     ggml_vk_create_pipeline(device, device->pipeline_flash_attn_split_k_reduce, "fa_split_k_reduce", fa_split_k_reduce_len, fa_split_k_reduce_data, "main", 3, sizeof(vk_op_flash_attn_split_k_reduce_push_constants), {1, device->subgroup_size, 1}, {device->subgroup_size}, 1, true);
 
 #if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
-    if (device->vendor_id == VK_VENDOR_ID_INTEL && (device->architecture == INTEL_XE2 || (device->architecture == INTEL_XE1 && device->coopmat_support && device->uma))) {
-        auto upper_power_of_2 = [&](uint32_t in) {
-            GGML_ASSERT(in != 0);
-            if (in <= 1) return 1u;
-            uint32_t ret = in - 1;
-            ret |= ret >> 1;
-            ret |= ret >> 2;
-            ret |= ret >> 4;
-            ret |= ret >> 8;
-            ret |= ret >> 16;
-            return ret + 1;
-            };
-
-        uint32_t xe_native_sub_group_size = 16;
-        if (device->architecture == INTEL_XE1) {
-            xe_native_sub_group_size = 8;
-        }
+    if (device->xe_fa_support) {
+        uint32_t shmem_sz = (uint32_t)device->properties.limits.maxComputeSharedMemorySize;
+        vk_xe_fa_tuning_params xe_fa_params {};
 
         for (auto& it : device->pipeline_xe_fa_decode_dual_phases) {
-            const uint32_t split_p_chunk = 32;
-            auto HdQk = it.first;
+            auto pipeline_state = it.first;
             auto& pipelines = it.second;
-            uint32_t head_dim_qk = std::get<0>(HdQk);
-            uint32_t head_dim_pv = std::get<1>(HdQk);
-            uint32_t gqa_ratio = std::get<2>(HdQk);
-            uint32_t q_len = std::get<3>(HdQk);
-            const uint32_t out_dim_per_wg = gqa_ratio > 16 ? 8 : 16;
-            uint32_t aligned_q_len = upper_power_of_2(q_len);
-            uint32_t group_sz_ph1 = std::min(std::max(aligned_q_len * xe_native_sub_group_size, 64u), 256u);
-            uint32_t out_per_wg_ph1 = std::min(q_len, 256u / xe_native_sub_group_size);
-            uint32_t aligned_gqa_ratio = upper_power_of_2(gqa_ratio);
-            uint32_t split_p_per_iter_ph2 = 256;
-            uint32_t split_p_per_warp = 16;
-            uint32_t group_sz_ph2 = (split_p_per_iter_ph2 / split_p_per_warp) * xe_native_sub_group_size;
-            uint32_t out_per_wg_ph2 = std::min(std::max(16u / aligned_gqa_ratio, 1u), q_len);
-            ggml_vk_create_pipeline(device, pipelines.first, "xe_fa_decode_ph1", fa_decode_ph1_cm1_len, fa_decode_ph1_cm1_data, "main", 5, sizeof(vk_fa_xe_opt_push_constants), { 1, 32, 1 }, { group_sz_ph1, gqa_ratio, head_dim_qk, xe_native_sub_group_size, split_p_chunk, out_per_wg_ph1 }, 1, false, true, xe_native_sub_group_size);
-            ggml_vk_create_pipeline(device, pipelines.second, "xe_fa_decode_ph2", fa_decode_ph2_cm1_len, fa_decode_ph2_cm1_data, "main", 5, sizeof(vk_fa_xe_opt_push_constants), { 1, 1, 1 }, { group_sz_ph2, gqa_ratio, head_dim_pv, out_per_wg_ph2, xe_native_sub_group_size, split_p_per_iter_ph2, split_p_chunk, out_dim_per_wg }, 1, false, true, xe_native_sub_group_size);
+            xe_fa_params = get_xe_fa_tuning_params(device.get(), pipeline_state, false);
+            std::vector<uint32_t> sc_ph1 = { xe_fa_params.workgroup_size[0], pipeline_state.qk_ratio, pipeline_state.HSK,
+                                             xe_fa_params.subgroup_size[0], xe_fa_params.decode_params.p_chunk_ph1,
+                                             xe_fa_params.decode_params.out_per_wg_ph1, xe_fa_params.decode_params.p_chunk_ph1 / xe_fa_params.tm,
+                                             xe_fa_params.tm, xe_fa_params.tn, shmem_sz };
+            std::vector<uint32_t> sc_ph2 = { xe_fa_params.workgroup_size[1], pipeline_state.pv_ratio, pipeline_state.HSV,
+                                             xe_fa_params.decode_params.out_row_per_wg_ph2, xe_fa_params.subgroup_size[1], xe_fa_params.decode_params.split_p_per_iter_ph2,
+                                             xe_fa_params.decode_params.p_chunk_ph1, xe_fa_params.decode_params.out_col_per_wg_ph2,
+                                             xe_fa_params.tm, xe_fa_params.tn, xe_fa_params.decode_params.pipeline_flag_ph2 };
+            ggml_vk_create_pipeline(device, pipelines.first, "xe_fa_decode_ph1", fa_decode_ph1_cm1_len, fa_decode_ph1_cm1_data, "main", 5, sizeof(vk_xe_fa_push_constants), { 1, xe_fa_params.decode_params.p_chunk_ph1, 1 }, sc_ph1, 1, false, true, xe_fa_params.subgroup_size[0]);
+            ggml_vk_create_pipeline(device, pipelines.second, "xe_fa_decode_ph2", fa_decode_ph2_cm1_len, fa_decode_ph2_cm1_data, "main", 5, sizeof(vk_xe_fa_push_constants), { 1, 1, 1 }, sc_ph2, 1, false, true, xe_fa_params.subgroup_size[0]);
+        }
+
+        for (auto& it : device->pipeline_xe_fa_prefill) {
+            auto pipeline_state = it.first;
+            auto& pipelines = it.second;
+            xe_fa_params = get_xe_fa_tuning_params(device.get(), pipeline_state, true);
+
+            std::vector<uint32_t> sc = { xe_fa_params.workgroup_size[0], xe_fa_params.prefill_params.wg_row, xe_fa_params.prefill_params.sg_col,
+                                         pipeline_state.HSK, pipeline_state.HSV, pipeline_state.qk_ratio, xe_fa_params.prefill_params.p_per_iter,
+                                         xe_fa_params.tm, xe_fa_params.tn, xe_fa_params.prefill_params.q_in_reg, shmem_sz, xe_fa_params.prefill_params.acc_count,
+                                         xe_fa_params.prefill_params.pipeline_flag, xe_fa_params.subgroup_size[0], xe_fa_params.prefill_params.gqa_fold };
+            ggml_vk_create_pipeline(device, pipelines, "xe_fa_prefill", fa_prefill_cm1_len, fa_prefill_cm1_data, "main", 6, sizeof(vk_xe_fa_push_constants), { xe_fa_params.prefill_params.gqa_fold, xe_fa_params.prefill_params.wg_row / xe_fa_params.prefill_params.gqa_fold, 1 }, sc, 1, false, true, xe_fa_params.subgroup_size[0]);
         }
     }
 #endif
@@ -3503,10 +3676,6 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     CREATE_UNARY_MUL(softplus, 3)
 #undef CREATE_UNARY_MUL
 
-    ggml_vk_create_pipeline(device, device->pipeline_add1_f16_f16, "add1_f16_f16", add1_f16_f16_len, add1_f16_f16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1);
-    ggml_vk_create_pipeline(device, device->pipeline_add1_f16_f32, "add1_f16_f32", add1_f16_f32_len, add1_f16_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1);
-    ggml_vk_create_pipeline(device, device->pipeline_add1_f32_f32, "add1_f32_f32", add1_f32_f32_len, add1_f32_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1);
-
     ggml_vk_create_pipeline(device, device->pipeline_arange_f32, "arange_f32", arange_f32_len, arange_f32_data, "main", 1, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);
 
     ggml_vk_create_pipeline(device, device->pipeline_fill_f32, "fill_f32", fill_f32_len, fill_f32_data, "main", 1, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);
@@ -3527,8 +3696,6 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
 
     ggml_vk_create_pipeline(device, device->pipeline_silu_back_f32, "silu_back_f32", silu_back_f32_len, silu_back_f32_data, "main", 3, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);
 
-    ggml_vk_create_pipeline(device, device->pipeline_diag_mask_inf_f32, "diag_mask_inf_f32", diag_mask_inf_f32_len, diag_mask_inf_f32_data, "main", 2, sizeof(vk_op_diag_mask_push_constants), {1, 512, 1}, {}, 1, true);
-
     ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32, "soft_max_f32", soft_max_f32_len, soft_max_f32_data, "main", 4, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
     ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32_wg512, "soft_max_f32_wg512", soft_max_f32_len, soft_max_f32_data, "main", 4, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { 512 }, 1);
     ggml_vk_create_pipeline(device, device->pipeline_soft_max_f32_f16, "soft_max_f32_f16", soft_max_f32_f16_len, soft_max_f32_f16_data, "main", 4, sizeof(vk_op_soft_max_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
@@ -4125,6 +4292,7 @@ vk_device ggml_vk_get_device(size_t idx) {
         bool dot2_f16_support = false;
         bool ocp_microscaling_extension = false;
         bool shader_float8_extension = false;
+        bool workgroup_memory_explicit_layout_support = false;
 
         for (const auto& properties : ext_props) {
             if (strcmp("VK_KHR_maintenance4", properties.extensionName) == 0) {
@@ -4180,6 +4348,10 @@ vk_device ggml_vk_get_device(size_t idx) {
                 dot2_f16_support = true;
             } else if (strcmp("VK_KHR_pipeline_executable_properties", properties.extensionName) == 0) {
                 pipeline_executable_properties_support = true;
+#if defined(VK_KHR_workgroup_memory_explicit_layout)
+            } else if (strcmp("VK_KHR_workgroup_memory_explicit_layout", properties.extensionName) == 0) {
+                workgroup_memory_explicit_layout_support = true;
+#endif
             } else if (strcmp("VK_EXT_memory_priority", properties.extensionName) == 0 &&
                        getenv("GGML_VK_ENABLE_MEMORY_PRIORITY")) {
                 device->memory_priority = true;
@@ -4532,6 +4704,16 @@ vk_device ggml_vk_get_device(size_t idx) {
             device_extensions.push_back("VK_KHR_pipeline_executable_properties");
         }
 
+#if defined(VK_KHR_workgroup_memory_explicit_layout)
+        VkPhysicalDeviceWorkgroupMemoryExplicitLayoutFeaturesKHR wg_explicit_layout_features {};
+        wg_explicit_layout_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_WORKGROUP_MEMORY_EXPLICIT_LAYOUT_FEATURES_KHR;
+        if (workgroup_memory_explicit_layout_support) {
+            last_struct->pNext = (VkBaseOutStructure *)&wg_explicit_layout_features;
+            last_struct = (VkBaseOutStructure *)&wg_explicit_layout_features;
+            device_extensions.push_back("VK_KHR_workgroup_memory_explicit_layout");
+        }
+#endif
+
         if (device->external_memory_host) {
             device_extensions.push_back("VK_EXT_external_memory_host");
         }
@@ -4593,6 +4775,16 @@ vk_device ggml_vk_get_device(size_t idx) {
                           shader_float8_extension && shader_float8_features.shaderFloat8 &&
                           !getenv("GGML_VK_DISABLE_OCP_FP4");
 
+#if defined(VK_KHR_workgroup_memory_explicit_layout)
+        device->workgroup_memory_explicit_layout = workgroup_memory_explicit_layout_support &&
+                                                   wg_explicit_layout_features.workgroupMemoryExplicitLayout;
+        device->workgroup_memory_explicit_layout_16bit = device->workgroup_memory_explicit_layout &&
+                                                         wg_explicit_layout_features.workgroupMemoryExplicitLayout16BitAccess;
+#else
+        device->workgroup_memory_explicit_layout = false;
+        device->workgroup_memory_explicit_layout_16bit = false;
+#endif
+
         device->pipeline_robustness = pl_robustness_features.pipelineRobustness;
 
         device->multi_add = vk12_props.shaderRoundingModeRTEFloat16 &&
@@ -4958,6 +5150,13 @@ vk_device ggml_vk_get_device(size_t idx) {
         descriptor_set_layout_create_info.setPNext(&dslbfci);
         device->dsl = device->device.createDescriptorSetLayout(descriptor_set_layout_create_info);
 
+#if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
+        device->xe_fa_support = device->coopmat_support && device->properties.limits.maxComputeSharedMemorySize >= 48 * 1024 &&
+                                device->workgroup_memory_explicit_layout && device->workgroup_memory_explicit_layout_16bit &&
+                                ((device->vendor_id == VK_VENDOR_ID_INTEL && (device->architecture == INTEL_XE2 || (device->architecture == INTEL_XE1 && device->uma))) ||
+                                 (device->subgroup_size == 32 && device->coopmat_support_16x16x16_f32acc && (getenv("GGML_VK_FORCE_XE_FA") != nullptr)));
+#endif
+
         ggml_vk_load_shaders(device);
 
         // Prefer a dedicated transfer queue on AMD dGPUs (non-GCN) when graphics queue use is disabled.
@@ -8403,32 +8602,98 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
 
     vk_pipeline pipeline = nullptr;
 
+    auto tensor_aligned_with = [&](const ggml_tensor * t, uint32_t in) {
+        GGML_ASSERT(t != nullptr && in != 0);
+        return t->nb[1] % in == 0 && t->nb[2] % in == 0 && t->nb[3] % in == 0 && ggml_vk_tensor_buffer_offset(ctx, t) % in == 0;
+    };
+    auto vk_dev = ctx->device.get();
     bool xe_fa_opt = false;
-    bool fa_copy_qstate = false;
-    bool xe_fa_supported_platform =
-        (ctx->device.get()->architecture == INTEL_XE2 && ctx->device.get()->properties.deviceID != 0xFD80 && ctx->device.get()->properties.deviceID != 0xFD81) ||
-        (ctx->device.get()->architecture == INTEL_XE1 && ctx->device.get()->coopmat_support && ctx->device.get()->uma);
-    bool xe_fa_supported_usage = neq0 % 32 == 0 && nev0 % 16 == 0 && q->nb[1] > q->nb[2] && k->nb[1] > k->nb[2] && v->nb[1] > v->nb[2] && mask != nullptr;
-    bool xe_fa_supported_dtype = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16 && (mask != nullptr && mask->type == GGML_TYPE_F16);
+    bool xe_fa_use_2phases = false;
+    bool xe_fa_supported_platform = vk_dev->xe_fa_support && vk_dev->properties.deviceID != 0xFD80 && vk_dev->properties.deviceID != 0xFD81;
+    bool xe_fa_supported_usage = neq0 % 32 == 0 && nev0 % 16 == 0 && mask != nullptr && !use_sparse && max_bias == 0.0f && logit_softcap == 0.0f &&
+                                 tensor_aligned_with(k, 16u) && tensor_aligned_with(v, 16u);
+    bool xe_fa_supported_dtype = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16 &&
+                                 mask != nullptr && mask->type == GGML_TYPE_F16;
+    xe_fa_opt = xe_fa_supported_platform && xe_fa_supported_usage && xe_fa_supported_dtype;
     std::pair<vk_pipeline, vk_pipeline> xe_fa_pipeline_dual_phases = { nullptr , nullptr };
     vk_pipeline xe_fa_pipeline = nullptr;
+    size_t size_q_staging = 0;
     size_t size_p = 0;
     size_t size_group_max = 0;
+    size_t size_decode_temp = 0;
+    uint32_t pv_ratio = uint32_t(neq2 / nev2);
+    vk_xe_fa_tuning_params xe_fa_params {};
 
     {
         std::lock_guard<std::mutex> guard(ctx->device->compile_mutex);
-        auto &pipelines = ctx->device->pipeline_flash_attn_f32_f16;
-        auto it = pipelines.find(fa_pipeline_state);
-        if (it != pipelines.end()) {
-            pipeline = it->second;
-        } else {
-            pipelines[fa_pipeline_state] = pipeline = std::make_shared<vk_pipeline_struct>();
+        if (xe_fa_opt) {
+            if (neq1 > 8 || pv_ratio > 64) {
+                uint32_t head_stride_q = uint32_t(nbq2 / ggml_type_size(q->type));
+                uint32_t batch_stride_q = uint32_t(nbq3 / ggml_type_size(q->type));
+                vk_xe_fa_pipeline_state prefill_state = { (uint32_t)neq0, (uint32_t)nev0, qk_ratio, pv_ratio, 0, false, false, false, false };
+                prefill_state.unaligned_seq_len = nek1 % 32 != 0;
+                prefill_state.unaligned_q = head_stride_q % 8 != 0 || q_stride % 8 != 0 || batch_stride_q % 8 != 0;
+                prefill_state.partial_row_mask = neq1 > nem1;
+                prefill_state.multi_head_mask = nem2 > 1;
+                xe_fa_params = get_xe_fa_tuning_params(vk_dev, prefill_state, true);
+                if (xe_fa_params.shmem_sz <= vk_dev->properties.limits.maxComputeSharedMemorySize) {
+                    auto& xe_fa_pipelines = ctx->device->pipeline_xe_fa_prefill;
+                    auto it = xe_fa_pipelines.find(prefill_state);
+                    if (it != xe_fa_pipelines.end()) {
+                        xe_fa_pipeline = it->second;
+                    } else {
+                        xe_fa_pipelines[prefill_state] = xe_fa_pipeline = std::make_shared<vk_pipeline_struct>();
+                    }
+                } else {
+                    xe_fa_opt = false;
+                }
+            } else {
+                vk_xe_fa_pipeline_state decode_state = { (uint32_t)neq0, (uint32_t)nev0, qk_ratio, pv_ratio, (uint32_t)neq1, false, false, false, false };
+                decode_state.unaligned_seq_len = nek1 % 256 != 0;
+                xe_fa_params = get_xe_fa_tuning_params(vk_dev, decode_state, false);
+                const uint32_t split_p_size = xe_fa_params.decode_params.p_chunk_ph1;
+                const size_t max_dim = std::max<size_t>(CEIL_DIV(nek1, split_p_size), 1);
+                const size_t p_dim = max_dim * split_p_size;
+                const size_t q_head_padding = 16;
+                size_q_staging = neq0 * neq1 * (neq2 * neq3 + q_head_padding) * sizeof(ggml_fp16_t);
+                size_p = neq1 * neq2 * p_dim * neq3 * sizeof(ggml_fp16_t);
+                size_group_max = neq1 * neq2 * max_dim * neq3 * sizeof(float);
+                size_decode_temp = size_q_staging + size_p + size_group_max;
+                uint32_t ph1_wg1 = CEIL_DIV(nek1, xe_fa_params.decode_params.p_chunk_ph1);
+                if (xe_fa_params.shmem_sz <= vk_dev->properties.limits.maxComputeSharedMemorySize &&
+                    ph1_wg1 <= ctx->device->properties.limits.maxComputeWorkGroupCount[1] &&
+                    size_decode_temp <= ctx->device->properties.limits.maxStorageBufferRange) {
+                    auto& pipelines = ctx->device->pipeline_xe_fa_decode_dual_phases;
+                    auto it = pipelines.find(decode_state);
+                    if (it != pipelines.end()) {
+                        xe_fa_pipeline_dual_phases = it->second;
+                    } else {
+                        pipelines[decode_state] = xe_fa_pipeline_dual_phases =
+                            std::make_pair(std::make_shared<vk_pipeline_struct>(), std::make_shared<vk_pipeline_struct>());
+                    }
+                    xe_fa_use_2phases = true;
+                } else {
+                    xe_fa_opt = false;
+                }
+            }
+        }
+
+        if (!xe_fa_opt) {
+            auto &pipelines = ctx->device->pipeline_flash_attn_f32_f16;
+            auto it = pipelines.find(fa_pipeline_state);
+            if (it != pipelines.end()) {
+                pipeline = it->second;
+            } else {
+                pipelines[fa_pipeline_state] = pipeline = std::make_shared<vk_pipeline_struct>();
+            }
         }
     }
 
-    assert(pipeline);
-    // Compile early to initialize wg_denoms.
-    ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
+    if (!xe_fa_opt) {
+        assert(pipeline);
+        // Compile early to initialize wg_denoms.
+        ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
+    }
 
     uint32_t split_kv = KV;
     uint32_t split_k = 1;
@@ -8442,7 +8707,9 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
     const uint32_t Br = fa_pipeline_state.Br;
     const uint32_t Bc = fa_pipeline_state.Bc;
 
-    GGML_ASSERT(Br == pipeline->wg_denoms[0]);
+    if (!xe_fa_opt) {
+        GGML_ASSERT(Br == pipeline->wg_denoms[0]);
+    }
     const uint32_t Tr = CEIL_DIV(N, Br);
 
     // Try to use split_k when KV is large enough to be worth the overhead.
@@ -8472,43 +8739,22 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
         // of "align", so recompute split_k based on that.
         split_kv = ROUNDUP_POW2(std::max(1u, KV / split_k), alignment);
         split_k = CEIL_DIV(KV, split_kv);
-        xe_fa_opt = xe_fa_supported_platform && xe_fa_supported_usage && xe_fa_supported_dtype;
-        if (xe_fa_opt) {
-            const uint32_t split_p_size = 32;
-            const size_t max_dim = (nek1 + split_p_size - 1) / split_p_size;
-            const size_t p_dim = max_dim * split_p_size;
-#if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
-            std::lock_guard<std::mutex> guard(ctx->device->compile_mutex);
-            auto& pipelines = ctx->device->pipeline_xe_fa_decode_dual_phases;
-            auto it = pipelines.find({ (uint32_t)neq0, (uint32_t)nev0, qk_ratio, (uint32_t)neq1 });
-            if (it != pipelines.end()) {
-                xe_fa_pipeline_dual_phases = it->second;
-            } else {
-                pipelines[{(uint32_t)neq0, (uint32_t)nev0, qk_ratio, (uint32_t)neq1}] = xe_fa_pipeline_dual_phases = std::make_pair(std::make_shared<vk_pipeline_struct>(), std::make_shared<vk_pipeline_struct>());
-            }
-#endif
-            if (xe_fa_pipeline_dual_phases.first == nullptr || xe_fa_pipeline_dual_phases.second == nullptr) {
-                xe_fa_opt = false;
-                fa_copy_qstate = false;
-            } else {
-                size_p = neq1 * neq2 * p_dim * neq3 * sizeof(ggml_fp16_t);
-                size_group_max = neq1 * neq2 * max_dim * neq3 * sizeof(float);
-                size_t temp_size = ggml_nelements(q) * sizeof(ggml_fp16_t) + size_p + size_group_max;
-                fa_copy_qstate = true;
-                if (ctx->prealloc_size_x < temp_size) {
-                    ctx->prealloc_size_x = temp_size;
-                    ggml_vk_preallocate_buffers(ctx, subctx);
-                }
-
-                if (ctx->prealloc_x_need_sync) {
-                    ggml_vk_sync_buffers(ctx, subctx);
-                }
-            }
-        }
     }
 
-    if (xe_fa_opt == true) {
+    if (xe_fa_opt) {
+        // Bypass mask_opt and split_k buffer grow when they are not in use
         use_mask_opt = false;
+        split_k = 1;
+        if (xe_fa_use_2phases) {
+            if (ctx->prealloc_size_x < size_decode_temp) {
+                ctx->prealloc_size_x = size_decode_temp;
+                ggml_vk_preallocate_buffers(ctx, subctx);
+            }
+
+            if (ctx->prealloc_x_need_sync) {
+                ggml_vk_sync_buffers(ctx, subctx);
+            }
+        }
     }
 
     // Reserve space for split_k temporaries. For each split x batch, we need to store the O matrix (D x ne1)
@@ -8697,67 +8943,59 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
                                               mask_n_head_log2, m0, m1,
                                               gqa_ratio, split_kv, split_k };
 
-    if (xe_fa_opt && split_k > 1) {
-        auto upper_power_of_2 = [&](uint32_t in) {
-            GGML_ASSERT(in != 0);
-            if (in <= 1) return 1u;
-            uint32_t ret = in - 1;
-            ret |= ret >> 1;
-            ret |= ret >> 2;
-            ret |= ret >> 4;
-            ret |= ret >> 8;
-            ret |= ret >> 16;
-            return ret + 1;
-        };
-        auto to_fp16_vk_0 = ggml_vk_get_to_fp16(ctx, q->type);
-        const uint32_t out_dim_per_wg = qk_ratio > 16 ? 8 : 16;
-        size_t x_ne = ggml_nelements(q);
-        size_t temp_buf_offset = 0;
+    if (xe_fa_opt) {
+        uint32_t head_stride_q = uint32_t(nbq2 / ggml_type_size(q->type));
         uint32_t head_stride_k = uint32_t(nbk2 / ggml_type_size(k->type));
         uint32_t head_stride_v = uint32_t(nbv2 / ggml_type_size(v->type));
         uint32_t batch_stride_q = uint32_t(nbq3 / ggml_type_size(q->type));
         uint32_t batch_stride_k = uint32_t(nbk3 / ggml_type_size(k->type));
         uint32_t batch_stride_v = uint32_t(nbv3 / ggml_type_size(v->type));
-        uint32_t batch_stride_m = mask ? uint32_t(mask->nb[3] / ggml_type_size(mask->type)) : 0u;
         uint32_t batch_stride_o = uint32_t(nb3 / ggml_type_size(dst->type));
-        vk_fa_xe_opt_push_constants pc_ph1 = { (uint32_t)nek1, (uint32_t)neq1, (uint32_t)neq2, (uint32_t)nek2, qk_ratio, 1, (sinks != nullptr) ? 1u : 0u, (uint32_t)k_stride, head_stride_k,
-            batch_stride_q, batch_stride_k, batch_stride_v, batch_stride_m, batch_stride_o, scale };
-        vk_fa_xe_opt_push_constants pc_ph2 = pc_ph1;
-        pc_ph2.nbkv_tok = v_stride;
-        pc_ph2.nbkv_head = head_stride_v;
-        vk_subbuffer q_temp_buf = fa_copy_qstate ? ggml_vk_subbuffer(ctx, ctx->prealloc_x, temp_buf_offset) : q_buf;
-        temp_buf_offset += fa_copy_qstate ? x_ne * sizeof(ggml_fp16_t) : 0;
-        vk_subbuffer p_temp_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_x, temp_buf_offset);
-        temp_buf_offset += size_p;
-        vk_subbuffer max_temp_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_x, temp_buf_offset);
-        temp_buf_offset += size_group_max;
-        uint32_t xe_native_sub_group_size = ctx->device.get()->architecture == INTEL_XE1 ? 8 : 16;
-        uint32_t aligned_gqa_ratio = upper_power_of_2(qk_ratio);
-        uint32_t out_per_wg_ph1 = std::min(256u / xe_native_sub_group_size, (uint32_t)neq1);
-        uint32_t out_per_wg_ph2 = std::min(std::max(16u / aligned_gqa_ratio, 1u), (uint32_t)neq1);
-        uint32_t ph1_wg = ((neq1 + out_per_wg_ph1 - 1) / out_per_wg_ph1) * nek2;
-        uint32_t ph2_wg = ((neq1 + out_per_wg_ph2 - 1) / out_per_wg_ph2) * ne0 / out_dim_per_wg;
-        if (fa_copy_qstate) {
-            const std::vector<uint32_t> pc_cpy_fp16 =
-            { (uint32_t)q->ne[0], (uint32_t)q->ne[1], (uint32_t)q->ne[2], (uint32_t)q->ne[3], (uint32_t)(x_ne) };
-            ggml_vk_sync_buffers(ctx, subctx);
+
+        vk_xe_fa_push_constants pc_xe_fa = { (uint32_t)neq1, (uint32_t)neq2, (uint32_t)neq3,
+                                             (uint32_t)nek1, (uint32_t)nek2, (uint32_t)nek3, (uint32_t)nev2, (uint32_t)nev3,
+                                             (uint32_t)nem1, (uint32_t)nem2, (uint32_t)nem3,
+                                             q_stride, head_stride_q, k_stride, head_stride_k, v_stride, head_stride_v,
+                                             batch_stride_q, batch_stride_k, batch_stride_v, batch_stride_o,
+                                             (sinks != nullptr) ? (1u << 16) : 0u, scale };
+
+        if (xe_fa_use_2phases) {
+            auto to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, q, nullptr, GGML_TYPE_F16);
+            size_t temp_buf_offset = 0;
+            vk_subbuffer q_temp_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_x, temp_buf_offset);
+            temp_buf_offset += size_q_staging;
+            vk_subbuffer p_temp_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_x, temp_buf_offset);
+            temp_buf_offset += size_p;
+            vk_subbuffer max_temp_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_x, temp_buf_offset);
+            temp_buf_offset += size_group_max;
+            uint32_t ph1_wg = CEIL_DIV(neq1, xe_fa_params.decode_params.out_per_wg_ph1) * nek2;
+            uint32_t ph2_wg = CEIL_DIV(neq1, xe_fa_params.decode_params.out_row_per_wg_ph2) * ne0 / xe_fa_params.decode_params.out_col_per_wg_ph2;
+
+            pc_xe_fa.nbq1 = uint32_t(neq0);
+            pc_xe_fa.nbq2 = uint32_t(neq0 * neq1);
+            pc_xe_fa.nbq3 = uint32_t(neq0 * neq1 * neq2);
+
             ggml_pipeline_request_descriptor_sets(ctx, to_fp16_vk_0, 1);
-            ggml_vk_dispatch_pipeline(ctx, subctx, to_fp16_vk_0, { q_buf, q_temp_buf }, pc_cpy_fp16, { (uint32_t)(x_ne), 1, 1 });
+            ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, q, q_buf, q_temp_buf);
+
+            ggml_pipeline_request_descriptor_sets(ctx, xe_fa_pipeline_dual_phases.first, 1);
+            ggml_vk_dispatch_pipeline(ctx, subctx, xe_fa_pipeline_dual_phases.first,
+                { q_temp_buf, k_buf, mask_buf, p_temp_buf, max_temp_buf },
+                pc_xe_fa, { (uint32_t)ph1_wg, (uint32_t)nek1, (uint32_t)neq3 });
+
+            ggml_vk_sync_buffers(ctx, subctx);
+            ggml_pipeline_request_descriptor_sets(ctx, xe_fa_pipeline_dual_phases.second, 1);
+            ggml_vk_dispatch_pipeline(ctx, subctx, xe_fa_pipeline_dual_phases.second,
+                { p_temp_buf, v_buf, max_temp_buf, sinks_buf, dst_buf },
+                pc_xe_fa, { (uint32_t)ph2_wg, (uint32_t)nev2, (uint32_t)neq3 });
+
+            ctx->prealloc_x_need_sync = true;
+        } else {
+            ggml_pipeline_request_descriptor_sets(ctx, xe_fa_pipeline, 1);
+            ggml_vk_dispatch_pipeline(ctx, subctx, xe_fa_pipeline,
+                { q_buf, k_buf, v_buf, mask_buf, sinks_buf, dst_buf },
+                pc_xe_fa, { (uint32_t)neq2, (uint32_t)neq1, (uint32_t)neq3 });
         }
-
-        ggml_vk_sync_buffers(ctx, subctx);
-        ggml_pipeline_request_descriptor_sets(ctx, xe_fa_pipeline_dual_phases.first, 1);
-        ggml_vk_dispatch_pipeline(ctx, subctx, xe_fa_pipeline_dual_phases.first,
-            { q_temp_buf, k_buf, mask_buf, p_temp_buf, max_temp_buf },
-            pc_ph1, { (uint32_t)ph1_wg, (uint32_t)nek1, (uint32_t)neq3 });
-
-        ggml_vk_sync_buffers(ctx, subctx);
-        ggml_pipeline_request_descriptor_sets(ctx, xe_fa_pipeline_dual_phases.second, 1);
-        ggml_vk_dispatch_pipeline(ctx, subctx, xe_fa_pipeline_dual_phases.second,
-            { p_temp_buf, v_buf, max_temp_buf, sinks_buf, dst_buf },
-            pc_ph2, { (uint32_t)ph2_wg, (uint32_t)nev2, (uint32_t)neq3 });
-
-        ctx->prealloc_x_need_sync = true;
     } else if (split_k > 1) {
         ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_split_k_reduce, 1);
 
@@ -9233,11 +9471,6 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
                 break;
         }
         return nullptr;
-    case GGML_OP_DIAG_MASK_INF:
-        if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
-            return ctx->device->pipeline_diag_mask_inf_f32;
-        }
-        return nullptr;
     case GGML_OP_SOFT_MAX:
         GGML_ASSERT(!src1 || src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16);
         GGML_ASSERT(!src2 || src2->type == GGML_TYPE_F32);
@@ -9616,17 +9849,6 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
             return pipeline;
         }
         return nullptr;
-    case GGML_OP_ADD1:
-        if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) {
-            return ctx->device->pipeline_add1_f16_f16;
-        }
-        if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) {
-            return ctx->device->pipeline_add1_f16_f32;
-        }
-        if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
-            return ctx->device->pipeline_add1_f32_f32;
-        }
-        return nullptr;
     case GGML_OP_ARANGE:
         if (dst->type == GGML_TYPE_F32) {
             return ctx->device->pipeline_arange_f32;
@@ -9765,9 +9987,6 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
             elements = { groups_x, groups_y, uint32_t(CEIL_DIV(uint64_t(total_workgroups), groups_xy)) };
             GGML_ASSERT(elements[2] <= ctx->device->properties.limits.maxComputeWorkGroupCount[2]);
         } break;
-    case GGML_OP_DIAG_MASK_INF:
-        elements = { (uint32_t)ggml_nrows(src0), (uint32_t)ne00, 1 };
-        break;
     case GGML_OP_ROPE:
     case GGML_OP_ROPE_BACK:
         {
@@ -9901,7 +10120,6 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
     case GGML_OP_SUB:
     case GGML_OP_DIV:
     case GGML_OP_MUL:
-    case GGML_OP_ADD1:
     case GGML_OP_OUT_PROD:
     case GGML_OP_ARANGE:
     case GGML_OP_FILL:
@@ -10796,21 +11014,6 @@ void ggml_vk_sqrt(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_
     ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_SQRT, vk_op_unary_push_constants_init(src0, dst));
 }
 
-void ggml_vk_add1(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
-    const uint32_t src0_type_size = ggml_type_size(src0->type);
-    const uint32_t src1_type_size = ggml_type_size(src1->type);
-    const uint32_t dst_type_size = ggml_type_size(dst->type);
-
-    ggml_vk_op_f32<vk_op_binary_push_constants>(ctx, subctx, src0, src1, nullptr, nullptr, dst, GGML_OP_ADD1, {
-        (uint32_t)ggml_nelements(src0),
-        (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size,
-        (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size,
-        (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] /  dst_type_size, (uint32_t) dst->nb[1] /  dst_type_size, (uint32_t) dst->nb[2] /  dst_type_size, (uint32_t) dst->nb[3] /  dst_type_size,
-        0,
-        0.0f, 0.0f, 0,
-    });
-}
-
 void ggml_vk_arange(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
     VK_LOG_DEBUG("ggml_vk_arange(dst=" << dst << ", ne=" << ggml_nelements(dst) << ")");
 
@@ -11325,11 +11528,6 @@ void ggml_vk_glu(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_t
         });
 }
 
-void ggml_vk_diag_mask_inf(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) {
-    int32_t * op_params = (int32_t *)dst->op_params;
-    ggml_vk_op_f32<vk_op_diag_mask_push_constants>(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_DIAG_MASK_INF, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], op_params[0] });
-}
-
 void ggml_vk_soft_max(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst) {
     float * op_params = (float *)dst->op_params;
 
@@ -12701,10 +12899,6 @@ bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, in
     case GGML_OP_UPSCALE:
         ggml_vk_upscale(ctx, compute_ctx, src0, node);
 
-        break;
-    case GGML_OP_ADD1:
-        ggml_vk_add1(ctx, compute_ctx, src0, src1, node);
-
         break;
     case GGML_OP_ARANGE:
         ggml_vk_arange(ctx, compute_ctx, node);
@@ -12855,10 +13049,6 @@ bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, in
         default:
             return false;
         }
-        break;
-    case GGML_OP_DIAG_MASK_INF:
-        ggml_vk_diag_mask_inf(ctx, compute_ctx, src0, node);
-
         break;
     case GGML_OP_SOFT_MAX:
         if (ctx->fused_topk_moe_mode != TOPK_MOE_COUNT) {
@@ -16125,10 +16315,6 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
         case GGML_OP_CONCAT: {
             return ggml_vk_concat_supported(op->src[0], op->src[1], op);
         }
-        case GGML_OP_ADD1:
-            return (op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32)
-                || (op->src[0]->type == GGML_TYPE_F16 && op->src[1]->type == GGML_TYPE_F32)
-                || (op->src[0]->type == GGML_TYPE_F16 && op->src[1]->type == GGML_TYPE_F16);
         case GGML_OP_ARANGE:
             return op->type == GGML_TYPE_F32;
         case GGML_OP_FILL:
@@ -16139,8 +16325,6 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
         case GGML_OP_PAD_REFLECT_1D:
         case GGML_OP_ROLL:
             return op->src[0]->type == GGML_TYPE_F32;
-        case GGML_OP_DIAG_MASK_INF:
-            return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_SOFT_MAX:
             return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32
                 && (!op->src[1] || (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16));

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

@@ -1,34 +0,0 @@
-#version 450
-
-#extension GL_EXT_shader_16bit_storage : require
-#extension GL_EXT_control_flow_attributes : enable
-
-layout (push_constant) uniform parameter
-{
-    uint ncols;
-    uint rows_per_channel;
-    uint n_past;
-} p;
-
-#include "types.glsl"
-
-layout(local_size_x = 1, local_size_y = 512, 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 col = gl_GlobalInvocationID.y;
-    const uint row = gl_GlobalInvocationID.x;
-
-    if (col >= p.ncols) {
-        return;
-    }
-
-    const uint i = row*p.ncols + col;
-    if (col > p.n_past + row % p.rows_per_channel) {
-        data_d[i] = D_TYPE(uintBitsToFloat(0xFF800000));
-    } else {
-        data_d[i] = D_TYPE(data_a[i]);
-    }
-}

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

@@ -3,261 +3,327 @@
 #extension GL_EXT_control_flow_attributes : enable
 #extension GL_EXT_shader_16bit_storage : require
 #extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
-#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
 #extension GL_KHR_memory_scope_semantics : enable
 #extension GL_KHR_shader_subgroup_basic : enable
-#extension GL_KHR_shader_subgroup_ballot : enable
 #extension GL_KHR_shader_subgroup_arithmetic : enable
 #extension GL_KHR_cooperative_matrix : enable
 #extension GL_EXT_shared_memory_block : enable
 
-layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+layout (local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
 
-layout (binding = 0) readonly buffer Q {float16_t qState[];};
-layout (binding = 1) readonly buffer K_VEC4 {f16vec4 kStateVec4[];};
-layout (binding = 2) buffer MASK_F16 {float16_t mState_f16[];};
-layout (binding = 3) buffer P_FP16 {float16_t matP_f16[];};
-layout (binding = 4) buffer OUT_MAX {float out_max_f32[];};
+layout (binding = 0) readonly buffer Q {float16_t data_q[];};
+layout (binding = 1) readonly buffer KV4 {f16vec4 data_kv4[];};
+layout (binding = 2) readonly buffer M {float16_t data_m[];};
+layout (binding = 3) buffer P {float16_t data_p[];};
+layout (binding = 4) buffer MAX {float data_max[];};
 
 layout (push_constant) uniform parameter
 {
-  uint kvSeqLen;
-  uint activationLength;
-  uint qHead;
-  uint kvHead;
-  uint qkRatio;
-  uint qkSubGroups;
-  uint flag;
-  uint kvStride1;
-  uint kvStride2;
-  uint batchStrideQ;
-  uint batchStrideK;
-  uint batchStrideV;
-  uint batchStrideM;
-  uint batchStrideO;
-  float softMaxScale;
+    uint neq1;
+    uint neq2;
+    uint neq3;
+    uint nek1;
+    uint nek2;
+    uint nek3;
+    uint nev2;
+    uint nev3;
+    uint nem1;
+    uint nem2;
+    uint nem3;
+    uint nbq1;
+    uint nbq2;
+    uint nbk1;
+    uint nbk2;
+    uint nbv1;
+    uint nbv2;
+    uint nbq3;
+    uint nbk3;
+    uint nbv3;
+    uint nb3;
+    uint flag;
+    float scale;
 } p;
 
-layout (constant_id = 0) const uint GROUPSIZE = 128;
-layout (constant_id = 1) const uint GQA_RATIO = 8;
-layout (constant_id = 2) const uint HEAD_DIM = 128;
-layout (constant_id = 3) const uint WARPSIZE = 16;
-layout (constant_id = 4) const uint MATP_REDUCE = 32;
-layout (constant_id = 5) const uint N_TOK = 1;
-layout (constant_id = 6) const uint COOP_MAT_P_PER_LOOP = 4;
+layout (constant_id = 0) const uint WorkGroupSize = 128;
+layout (constant_id = 1) const uint gqa_ratio = 8;
+layout (constant_id = 2) const uint HSK = 128;
+layout (constant_id = 3) const uint SubGroupSize = 16;
+layout (constant_id = 4) const uint p_reduce = 32;
+layout (constant_id = 5) const uint Br = 1;
+layout (constant_id = 6) const uint pmat_count = 4;
+layout (constant_id = 7) const uint TM = 8;
+layout (constant_id = 8) const uint TN = 16;
+layout (constant_id = 9) const uint shmem_size = 49152;
 
-#define MAX_HEADS 8
+const uint TK = 16;
+const uint sg_count = WorkGroupSize / SubGroupSize;
+const uint p_per_loop = pmat_count * TM;
+const uint p_loops = p_reduce / p_per_loop;
 
-#define TN WARPSIZE
-#define TM 8
-#define TK 16
-#define SUBGROUP_COUNT (GROUPSIZE / WARPSIZE)
-#define MATP_PER_LOOP (COOP_MAT_P_PER_LOOP * TM)
-#define P_LOOP_COUNT (MATP_REDUCE / MATP_PER_LOOP)
+const uint pmat_row = (gqa_ratio + TN - 1) / TN;
+const uint pmat_col = p_reduce / TM;
 
-#define COOP_MAT_Q_PER_TOKEN ((GQA_RATIO + TN - 1) / TN)
-#define COOP_MAT_P_M COOP_MAT_Q_PER_TOKEN
-#define COOP_MAT_P_N (MATP_REDUCE / TM)
-#define SLM_PV_SIZE (MATP_REDUCE * COOP_MAT_P_M * TN)
-#define SLM_MASK_SIZE (N_TOK * MATP_REDUCE)
-#define SLM_POOL_SIZE_K (MATP_PER_LOOP * HEAD_DIM)
-#define K_LOAD_PER_LOOP (GROUPSIZE * 4)
-#define HEAD_DIM_VEC4 (HEAD_DIM / 4)
-#define SLM_CHUNK_SIZE (TK / 4)
-#define K_LOAD_LOOPS ((SLM_POOL_SIZE_K + K_LOAD_PER_LOOP - 1) / K_LOAD_PER_LOOP)
-#define O_COUNT ((GQA_RATIO + SUBGROUP_COUNT - 1) / SUBGROUP_COUNT)
+const uint k_per_loop = WorkGroupSize * 4;
+const uint hsk_v4 = HSK / 4;
+const uint slm_col_chunk = TK / 4;
+const uint k_loops = (p_per_loop * HSK + k_per_loop - 1) / k_per_loop;
+const uint warp_o_row = (gqa_ratio + sg_count - 1) / sg_count;
 
-shared slm_pool_block {
-  float slm_pool_pv[SLM_PV_SIZE + SLM_MASK_SIZE];
-} slm_pool_f32;
+#define SLM_PV_SIZE (p_reduce * pmat_row * TN)
+#define SLM_MASK_SIZE (Br * p_reduce)
 
-shared slm_pool_alias_block {
-  float16_t slm_pool_k[SLM_POOL_SIZE_K];
-} slm_pool_f16;
+const bool k_in_shmem = k_loops * k_per_loop <= shmem_size / 2;
+const uint ksh_size = k_in_shmem ? (k_loops * k_per_loop) : (sg_count * TM * TK);
+
+shared sh_block {
+    float F32sh[SLM_PV_SIZE + SLM_MASK_SIZE];
+} F32Poolsh;
+
+shared shf16_alias {
+    float16_t F16sh[ksh_size];
+} F16Poolsh;
+
+shared shf16v4_alias {
+    f16vec4 F16V4sh[ksh_size / 4];
+} F16V4Poolsh;
 
 void main() {
-  const uint lane = gl_SubgroupInvocationID;
-  const uint kHeadIdx = gl_WorkGroupID.x % p.kvHead;
-  const uint outGroupIdx = gl_WorkGroupID.x / p.kvHead;
-  const uint v = gl_WorkGroupID.y;
-  const uint d = gl_WorkGroupID.z;
-  const uint localLinearId = gl_SubgroupID;
-  const uint wgLane = localLinearId * WARPSIZE + lane;
-  const uint qDim = p.qHead * HEAD_DIM;
-  const uint kvDim = p.kvStride1;
-  const uint maskDim = p.kvSeqLen;
-  const uint maxDim = (p.kvSeqLen + MATP_REDUCE - 1) / MATP_REDUCE;
-  const uint pDim = maxDim * MATP_REDUCE;
-  const uint tokFlatIdx = localLinearId + outGroupIdx * N_TOK;
-  uint offsetBaseQ = min(tokFlatIdx, p.activationLength - 1) * qDim;
-  offsetBaseQ = offsetBaseQ + d * p.batchStrideQ + kHeadIdx * HEAD_DIM * GQA_RATIO;
-  const uint offsetBaseK = (d * p.batchStrideK + (v * MATP_REDUCE) * kvDim + kHeadIdx * p.kvStride2) / 4;
-  uint offsetOut = d * p.qHead * p.activationLength * pDim + v * MATP_REDUCE + kHeadIdx * GQA_RATIO * pDim + (localLinearId * O_COUNT + outGroupIdx * N_TOK * p.qHead) * pDim + lane;
-  uint offsetMax = d * p.qHead * p.activationLength * maxDim + v + kHeadIdx * GQA_RATIO * maxDim + (localLinearId * O_COUNT + outGroupIdx * N_TOK * p.qHead) * maxDim;
-  const uint offsetSlmLoadPv = (localLinearId * O_COUNT * MATP_REDUCE + lane);
-  const uint offsetBaseM = v * MATP_REDUCE + lane;
-  const float fp32Min = uintBitsToFloat(0xFEFFFFFF);
+    const uint lane = gl_SubgroupInvocationID;
+    const uint head_k_idx = gl_WorkGroupID.x % p.nek2;
+    const uint out_idx = gl_WorkGroupID.x / p.nek2;
+    const uint p_idx = gl_WorkGroupID.y;
+    const uint batch_k_ratio = p.neq3 / p.nek3;
+    const uint batch_idx = gl_WorkGroupID.z;
+    const uint batch_k_idx = batch_idx / batch_k_ratio;
+    const uint batch_mask_idx = batch_idx % p.nem3;
+    const uint batch_mask_stride = p.nek1 * p.nem1 * p.nem2;
+    const uint sg = gl_SubgroupID;
+    const uint wg_lane = sg * SubGroupSize + lane;
+    const uint k_dim = p.nbk1;
+    const uint max_dim = (p.nek1 + p_reduce - 1) / p_reduce;
+    const uint p_dim = max_dim * p_reduce;
+    const uint out_row_idx = sg + out_idx * Br;
+    const bool single_head_mask = p.nem2 == 1;
+    uint q_offset = min(out_row_idx, p.neq1 - 1) * p.nbq1;
+    q_offset = q_offset + batch_idx * p.nbq3 + head_k_idx * p.nbq2 * gqa_ratio;
+    const uint k_base = (batch_k_idx * p.nbk3 + (p_idx * p_reduce) * k_dim + head_k_idx * p.nbk2) / 4;
+    uint o_offset = batch_idx * p.neq2 * p.neq1 * p_dim + p_idx * p_reduce + head_k_idx * gqa_ratio * p_dim +
+                    (sg * warp_o_row + out_idx * Br * p.neq2) * p_dim + lane;
+    uint max_offset = batch_idx * p.neq2 * p.neq1 * max_dim + p_idx + head_k_idx * gqa_ratio * max_dim +
+                      (sg * warp_o_row + out_idx * Br * p.neq2) * max_dim;
+    const uint pvsh_load_offset = sg * warp_o_row * p_reduce + lane;
+    const uint mask_base = p_idx * p_reduce + lane;
+    const uint k_seq_boundary = ((p.nek1 - 1) * k_dim + batch_k_idx * p.nbk3 + head_k_idx * p.nbk2 + HSK) / 4 - 1;
+    const uint k_head_base = (batch_k_idx * p.nbk3 + head_k_idx * p.nbk2) / 4;
+    const float NEG_FLT_MAX_OVER_2 = uintBitsToFloat(0xFEFFFFFF);
 
-  const uint loopCount = HEAD_DIM / TK;
-  float maskFp32[MATP_REDUCE / WARPSIZE];
+    float Maskf[p_reduce / SubGroupSize];
 
-  if (tokFlatIdx < p.activationLength) {
-    [[unroll]] for (uint mk = 0; mk < MATP_REDUCE / WARPSIZE; mk++) {
-      const uint maskOffset = mk * WARPSIZE + offsetBaseM;
-      if (maskOffset < maskDim) {
-        maskFp32[mk] = float(mState_f16[d * p.batchStrideM + tokFlatIdx * maskDim + maskOffset]);
-      } else {
-        maskFp32[mk] = fp32Min;
-      }
-    }
-  }
-
-  coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> matP[COOP_MAT_P_M][COOP_MAT_P_N];
-
-  [[unroll]] for (uint mp = 0; mp < COOP_MAT_P_M; mp++) {
-    [[unroll]] for (uint np = 0; np < COOP_MAT_P_N; np++) {
-      matP[mp][np] = coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(0.0f);
-    }
-  }
-
-  [[unroll]] for (uint kLoad = 0; kLoad < K_LOAD_LOOPS; kLoad++) {
-    const uint flatOffset = kLoad * GROUPSIZE + wgLane;
-    const uint kRowIdx = flatOffset / HEAD_DIM_VEC4;
-    const uint kColIdx = flatOffset % HEAD_DIM_VEC4;
-    const uint slmChunkCol = kColIdx % SLM_CHUNK_SIZE;
-    const uint slmChunkRow = kColIdx / SLM_CHUNK_SIZE;
-    const uint offsetK = offsetBaseK + kRowIdx * kvDim / 4 + kColIdx;
-    const uint offsetSlmK = kRowIdx * TK + slmChunkRow * TK * MATP_PER_LOOP + slmChunkCol * 4;
-    slm_pool_f16.slm_pool_k[offsetSlmK + 0] = kStateVec4[offsetK].x;
-    slm_pool_f16.slm_pool_k[offsetSlmK + 1] = kStateVec4[offsetK].y;
-    slm_pool_f16.slm_pool_k[offsetSlmK + 2] = kStateVec4[offsetK].z;
-    slm_pool_f16.slm_pool_k[offsetSlmK + 3] = kStateVec4[offsetK].w;
-  }
-
-  [[unroll]] for (uint pLoop = 0; pLoop < P_LOOP_COUNT; pLoop++) {
-    f16vec4 kTemp[K_LOAD_LOOPS];
-
-    if (pLoop + 1 < P_LOOP_COUNT) {
-      [[unroll]] for (uint kLoad = 0; kLoad < K_LOAD_LOOPS; kLoad++) {
-        const uint flatOffset = kLoad * GROUPSIZE + wgLane;
-        const uint kRowIdx = flatOffset / HEAD_DIM_VEC4 + (pLoop + 1) * MATP_PER_LOOP;
-        const uint kColIdx = flatOffset % HEAD_DIM_VEC4;
-        const uint offsetK = offsetBaseK + kRowIdx * kvDim / 4 + kColIdx;
-        kTemp[kLoad] = kStateVec4[offsetK];
-      }
+    if (out_row_idx < p.neq1 && single_head_mask) {
+        if (out_row_idx < p.nem1) {
+            [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+                const uint mask_offset = idx * SubGroupSize + mask_base;
+                if (mask_offset < p.nek1) {
+                    Maskf[idx] = float(data_m[batch_mask_idx * batch_mask_stride + out_row_idx * p.nek1 + mask_offset]);
+                } else {
+                    Maskf[idx] = NEG_FLT_MAX_OVER_2;
+                }
+            }
+        } else {
+            [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+                const uint mask_offset = idx * SubGroupSize + mask_base;
+                Maskf[idx] = (mask_offset < p.nek1) ? 0.0f : NEG_FLT_MAX_OVER_2;
+            }
+        }
     }
 
-    barrier();
-    if (localLinearId < N_TOK) {
-      [[unroll]] for (uint loop = 0; loop < loopCount; loop++) {
-        coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> matQ[COOP_MAT_P_M];
-        coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> matK[COOP_MAT_P_PER_LOOP];
+    coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> PMat[pmat_row][pmat_col];
 
-        [[unroll]] for (uint mq = 0; mq < COOP_MAT_P_M; mq++) {
-          coopMatLoad(
-            matQ[mq],
-            qState,
-            offsetBaseQ + mq * TN * HEAD_DIM + loop * TK,
-            HEAD_DIM,
-            gl_CooperativeMatrixLayoutColumnMajor);
+    [[unroll]] for (uint r = 0; r < pmat_row; ++r) {
+        [[unroll]] for (uint c = 0; c < pmat_col; ++c) {
+            PMat[r][c] = coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(0.0f);
+        }
+    }
+
+    if (k_in_shmem) {
+        [[unroll]] for (uint idx = 0; idx < k_loops; ++idx) {
+            const uint flat_offset = idx * WorkGroupSize + wg_lane;
+            const uint row_idx = flat_offset / hsk_v4;
+            const uint col_idx = flat_offset % hsk_v4;
+            const uint sh_chunk_c = col_idx % slm_col_chunk;
+            const uint sh_chunk_r = col_idx / slm_col_chunk;
+            uint k_offset = k_base + row_idx * k_dim / 4 + col_idx;
+            k_offset = min(k_offset, k_seq_boundary);
+            const uint ksh_offset = (row_idx * TK + sh_chunk_r * TK * p_per_loop) / 4 + sh_chunk_c;
+            F16V4Poolsh.F16V4sh[ksh_offset] = data_kv4[k_offset];
+        }
+    }
+
+    [[unroll]] for (uint pl = 0; pl < p_loops; ++pl) {
+        f16vec4 kv4_temp[k_loops];
+
+        if (k_in_shmem) {
+            if (pl + 1 < p_loops) {
+                [[unroll]] for (uint idx = 0; idx < k_loops; ++idx) {
+                    const uint flat_offset = idx * WorkGroupSize + wg_lane;
+                    const uint row_idx = flat_offset / hsk_v4 + (pl + 1) * p_per_loop;
+                    const uint col_idx = flat_offset % hsk_v4;
+                    uint k_offset = k_base + row_idx * k_dim / 4 + col_idx;
+                    k_offset = min(k_offset, k_seq_boundary);
+                    kv4_temp[idx] = data_kv4[k_offset];
+                }
+            }
         }
 
-        [[unroll]] for (uint np = 0; np < COOP_MAT_P_PER_LOOP; np++) {
-          coopMatLoad(
-            matK[np],
-            slm_pool_f16.slm_pool_k,
-            loop * TK * MATP_PER_LOOP + np * TM * TK,
-            TK,
-            gl_CooperativeMatrixLayoutRowMajor);
+        barrier();
+        if (sg < Br) {
+            [[unroll]] for (uint loop = 0; loop < HSK / TK; ++loop) {
+                coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> QMat[pmat_row];
+                coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> KMat[pmat_count];
+
+                [[unroll]] for (uint r = 0; r < pmat_row; ++r) {
+                    const uint coord = q_offset + r * TN * p.nbq2 + loop * TK;
+                    coopMatLoad(QMat[r], data_q, coord, p.nbq2, gl_CooperativeMatrixLayoutColumnMajor);
+                }
+
+                [[unroll]] for (uint c = 0; c < pmat_count; ++c) {
+                    if (k_in_shmem) {
+                        const uint coord = loop * TK * p_per_loop + c * TM * TK;
+                        coopMatLoad(KMat[c], F16Poolsh.F16sh, coord, TK, gl_CooperativeMatrixLayoutRowMajor);
+                    } else {
+                        const uint k_base_row = (c * TM + pl * pmat_count * TM + p_idx * p_reduce);
+                        const uint kv4_base_offset = k_head_base + k_base_row * p.nbk1 / 4 + loop * TK / 4;
+                        [[unroll]] for (uint idx = 0; idx < TM * TK / 4 / SubGroupSize; ++idx) {
+                            const uint flat_offset = idx * SubGroupSize + lane;
+                            const uint row_idx = flat_offset / (TK / 4);
+                            const uint col_idx = flat_offset % (TK / 4);
+                            const uint k_temp_offset = kv4_base_offset + row_idx * p.nbk1 / 4 + col_idx;
+                            const uint ksh_temp_offset = sg * TM * TK / 4 + idx * SubGroupSize + lane;
+                            f16vec4 kv4_temp = (k_base_row + row_idx) < p.nek1 ? data_kv4[k_temp_offset] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+                            F16V4Poolsh.F16V4sh[ksh_temp_offset] = kv4_temp;
+                        }
+
+                        subgroupBarrier();
+                        coopMatLoad(KMat[c], F16Poolsh.F16sh, sg * TM * TK, TK, gl_CooperativeMatrixLayoutRowMajor);
+                        subgroupBarrier();
+                    }
+                }
+
+                [[unroll]] for (uint r = 0; r < pmat_row; ++r) {
+                    [[unroll]] for (uint c = 0; c < pmat_count; ++c) {
+                        PMat[r][pl * pmat_count + c] = coopMatMulAdd(KMat[c], QMat[r], PMat[r][pl * pmat_count + c]);
+                    }
+                }
+            }
         }
 
-        [[unroll]] for (uint mp = 0; mp < COOP_MAT_P_M; mp++) {
-          [[unroll]] for (uint np = 0; np < COOP_MAT_P_PER_LOOP; np++) {
-            matP[mp][pLoop * COOP_MAT_P_PER_LOOP + np] = coopMatMulAdd(matK[np], matQ[mp], matP[mp][pLoop * COOP_MAT_P_PER_LOOP + np]);
-          }
+        barrier();
+
+        if (k_in_shmem) {
+            if (pl + 1 < p_loops) {
+                [[unroll]] for (uint idx = 0; idx < k_loops; ++idx) {
+                    const uint flat_offset = idx * WorkGroupSize + wg_lane;
+                    const uint row_idx = flat_offset / hsk_v4;
+                    const uint col_idx = flat_offset % hsk_v4;
+                    const uint sh_chunk_c = col_idx % slm_col_chunk;
+                    const uint sh_chunk_r = col_idx / slm_col_chunk;
+                    const uint ksh_offset = (row_idx * TK + sh_chunk_r * TK * p_per_loop) / 4 + sh_chunk_c;
+                    F16V4Poolsh.F16V4sh[ksh_offset] = kv4_temp[idx];
+                }
+            }
         }
-      }
     }
 
     barrier();
 
-    if (pLoop + 1 < P_LOOP_COUNT) {
-      [[unroll]] for (uint kLoad = 0; kLoad < K_LOAD_LOOPS; kLoad++) {
-        const uint flatOffset = kLoad * GROUPSIZE + wgLane;
-        const uint kRowIdx = flatOffset / HEAD_DIM_VEC4;
-        const uint kColIdx = flatOffset % HEAD_DIM_VEC4;
-        const uint slmChunkCol = kColIdx % SLM_CHUNK_SIZE;
-        const uint slmChunkRow = kColIdx / SLM_CHUNK_SIZE;
-        const uint offsetSlmK = kRowIdx * TK + slmChunkRow * TK * MATP_PER_LOOP + slmChunkCol * 4;
-        slm_pool_f16.slm_pool_k[offsetSlmK + 0] = kTemp[kLoad].x;
-        slm_pool_f16.slm_pool_k[offsetSlmK + 1] = kTemp[kLoad].y;
-        slm_pool_f16.slm_pool_k[offsetSlmK + 2] = kTemp[kLoad].z;
-        slm_pool_f16.slm_pool_k[offsetSlmK + 3] = kTemp[kLoad].w;
-      }
+    if (sg < Br && out_row_idx < p.neq1 && single_head_mask) {
+        [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+            F32Poolsh.F32sh[SLM_PV_SIZE + sg * p_reduce + idx * SubGroupSize + lane] = Maskf[idx];
+        }
     }
-  }
 
-  barrier();
+    [[unroll]] for (uint out_r = 0; out_r < Br; ++out_r) {
+        if (out_r + out_idx * Br < p.neq1) {
+            const uint mask_global_base = batch_mask_idx * batch_mask_stride + (out_r + out_idx * Br) * p.nek1;
+            if (sg == out_r) {
+                [[unroll]] for (uint r = 0; r < pmat_row; ++r) {
+                    [[unroll]] for (uint c = 0; c < pmat_col; ++c) {
+                        const uint coord = r * p_reduce * TN + c * TM;
+                        coopMatStore(PMat[r][c], F32Poolsh.F32sh, coord, p_reduce, gl_CooperativeMatrixLayoutColumnMajor);
+                    }
+                }
+            }
 
-  if (tokFlatIdx < p.activationLength) {
-    [[unroll]] for (uint mk = 0; mk < MATP_REDUCE / WARPSIZE; mk++) {
-      slm_pool_f32.slm_pool_pv[SLM_PV_SIZE + localLinearId * MATP_REDUCE + mk * WARPSIZE + lane] = maskFp32[mk];
+            barrier();
+
+            if (single_head_mask) {
+                [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+                    Maskf[idx] = F32Poolsh.F32sh[SLM_PV_SIZE + out_r * p_reduce + idx * SubGroupSize + lane];
+                }
+            }
+
+            float Of[warp_o_row][p_reduce / SubGroupSize];
+            float Mf[warp_o_row];
+
+            [[unroll]] for (uint r = 0; r < warp_o_row; ++r) {
+                const uint head_q_idx = head_k_idx * gqa_ratio + sg * warp_o_row + r;
+                const uint mask_head_offset = (head_q_idx % p.nem2) * p.nek1 * p.nem1;
+
+                [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+                    Of[r][idx] = F32Poolsh.F32sh[pvsh_load_offset + idx * SubGroupSize + r * p_reduce] * p.scale;
+                }
+
+                if (!single_head_mask) {
+                    if (out_r + out_idx * Br < p.nem1) {
+                        [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+                            const uint mask_offset = idx * SubGroupSize + mask_base;
+                            if (mask_offset < p.nek1) {
+                                Maskf[idx] = float(data_m[mask_global_base + mask_head_offset + mask_offset]);
+                            } else {
+                                Maskf[idx] = NEG_FLT_MAX_OVER_2;
+                            }
+                        }
+                    } else {
+                        [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+                            const uint mask_offset = idx * SubGroupSize + mask_base;
+                            Maskf[idx] = (mask_offset < p.nek1) ? 0.0f : NEG_FLT_MAX_OVER_2;
+                        }
+                    }
+                }
+
+                [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+                    Of[r][idx] = Of[r][idx] + Maskf[idx];
+                }
+
+                float Mtempf = NEG_FLT_MAX_OVER_2;
+                [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+                    const uint p_col_offset = idx * SubGroupSize + mask_base;
+                    if (p_col_offset < p.nek1) {
+                        Mtempf = max(Mtempf, Of[r][idx]);
+                    }
+                }
+                Mf[r] = subgroupMax(Mtempf);
+                [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+                    const uint p_col_offset = idx * SubGroupSize + mask_base;
+                    Of[r][idx] = (p_col_offset < p.nek1) ? exp(Of[r][idx] - Mf[r]) : 0.0f;
+                }
+            }
+
+            [[unroll]] for (uint r = 0; r < warp_o_row; ++r) {
+                if (sg * warp_o_row + r < gqa_ratio) {
+                    [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+                        data_p[o_offset + r * p_dim + idx * SubGroupSize] = float16_t(Of[r][idx]);
+                    }
+
+                    if (lane == 0) {
+                        data_max[max_offset + r * max_dim] = Mf[r];
+                    }
+                }
+            }
+
+            o_offset = o_offset + p.neq2 * p_dim;
+            max_offset = max_offset + p.neq2 * max_dim;
+            barrier();
+        }
     }
-  }
-
-  [[unroll]] for (uint oLoop = 0; oLoop < N_TOK; oLoop++) {
-    if (oLoop + outGroupIdx * N_TOK < p.activationLength) {
-      if (localLinearId == oLoop) {
-        [[unroll]] for (uint mp = 0; mp < COOP_MAT_P_M; mp++) {
-          [[unroll]] for (uint np = 0; np < COOP_MAT_P_N; np++) {
-            coopMatStore(matP[mp][np], slm_pool_f32.slm_pool_pv, mp * MATP_REDUCE * TN + np * TM, MATP_REDUCE, gl_CooperativeMatrixLayoutColumnMajor);
-          }
-        }
-      }
-
-      barrier();
-
-      [[unroll]] for (uint maskIdx = 0; maskIdx < MATP_REDUCE / WARPSIZE; maskIdx++) {
-        maskFp32[maskIdx] = slm_pool_f32.slm_pool_pv[SLM_PV_SIZE + oLoop * MATP_REDUCE + maskIdx * WARPSIZE + lane];
-      }
-
-      float fp32O[O_COUNT][MATP_REDUCE / WARPSIZE];
-      float maxOut[O_COUNT];
-
-      [[unroll]] for (uint oc = 0; oc < O_COUNT; oc++) {
-        [[unroll]] for (uint os = 0; os < MATP_REDUCE / WARPSIZE; os++) {
-          fp32O[oc][os] = slm_pool_f32.slm_pool_pv[offsetSlmLoadPv + os * WARPSIZE + oc * MATP_REDUCE] * p.softMaxScale;
-        }
-
-        [[unroll]] for (uint os = 0; os < MATP_REDUCE / WARPSIZE; os++) {
-          fp32O[oc][os] = fp32O[oc][os] + maskFp32[os];
-        }
-
-        float maxTemp = fp32Min;
-        [[unroll]] for (uint os = 0; os < MATP_REDUCE / WARPSIZE; os++) {
-          maxTemp = max(maxTemp, fp32O[oc][os]);
-        }
-        maxOut[oc] = subgroupMax(maxTemp);
-        [[unroll]] for (uint os = 0; os < MATP_REDUCE / WARPSIZE; os++) {
-          fp32O[oc][os] = exp(fp32O[oc][os] - maxOut[oc]);
-        }
-      }
-
-      [[unroll]] for (uint oc = 0; oc < O_COUNT; oc++) {
-        if (localLinearId * O_COUNT + oc < GQA_RATIO) {
-          [[unroll]] for (uint os = 0; os < MATP_REDUCE / WARPSIZE; os++) {
-            matP_f16[offsetOut + oc * pDim + os * WARPSIZE] = float16_t(fp32O[oc][os]);
-          }
-
-          if (lane == 0) {
-            out_max_f32[offsetMax + oc * maxDim] = maxOut[oc];
-          }
-        }
-      }
-
-      offsetOut = offsetOut + p.qHead * pDim;
-      offsetMax = offsetMax + p.qHead * maxDim;
-      barrier();
-    }
-  }
 }

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

@@ -3,406 +3,470 @@
 #extension GL_EXT_control_flow_attributes : enable
 #extension GL_EXT_shader_16bit_storage : require
 #extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
-#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
 #extension GL_KHR_memory_scope_semantics : enable
 #extension GL_KHR_shader_subgroup_basic : enable
-#extension GL_KHR_shader_subgroup_ballot : enable
 #extension GL_KHR_shader_subgroup_arithmetic : enable
 #extension GL_KHR_cooperative_matrix : enable
 #extension GL_EXT_shared_memory_block : enable
 
-layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+layout (local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
 
-layout (binding = 0) readonly buffer P {f16vec4 pStateVec4[];};
-layout (binding = 1) readonly buffer V {float16_t vState[];};
-layout (binding = 1) readonly buffer V_VEC4 {f16vec4 vStateVec4[];};
-layout (binding = 2) buffer MAX_FP32 {float max_f32[];};
-layout (binding = 3) buffer SINK_FP32 {float sink_f32[];};
-layout (binding = 4) buffer OUT_FP32 {float out_f32[];};
-layout (binding = 4) buffer OUT_VEC4 {vec4 out_f32_vec4[];};
-layout (binding = 4) buffer OUT_F16 {float16_t out_f16[];};
+layout (binding = 0) readonly buffer PV4 {f16vec4 data_pv4[];};
+layout (binding = 1) readonly buffer V {float16_t data_v[];};
+layout (binding = 1) readonly buffer VV4 {f16vec4 data_vv4[];};
+layout (binding = 2) readonly buffer MAX {float data_max[];};
+layout (binding = 3) readonly buffer S {float data_s[];};
+layout (binding = 4) buffer OV4 {vec4 data_ov4[];};
 
 layout (push_constant) uniform parameter
 {
-  uint kvSeqLen;
-  uint activationLength;
-  uint qHead;
-  uint kvHead;
-  uint qkRatio;
-  uint qkSubGroups;
-  uint flag;
-  uint kvStride1;
-  uint kvStride2;
-  uint batchStrideQ;
-  uint batchStrideK;
-  uint batchStrideV;
-  uint batchStrideM;
-  uint batchStrideO;
-  float softMaxScale;
+    uint neq1;
+    uint neq2;
+    uint neq3;
+    uint nek1;
+    uint nek2;
+    uint nek3;
+    uint nev2;
+    uint nev3;
+    uint nem1;
+    uint nem2;
+    uint nem3;
+    uint nbq1;
+    uint nbq2;
+    uint nbk1;
+    uint nbk2;
+    uint nbv1;
+    uint nbv2;
+    uint nbq3;
+    uint nbk3;
+    uint nbv3;
+    uint nb3;
+    uint flag;
+    float scale;
 } p;
 
-layout (constant_id = 0) const uint GROUPSIZE = 256;
-layout (constant_id = 1) const uint GQA_RATIO = 8;
-layout (constant_id = 2) const uint HEAD_DIM = 128;
-layout (constant_id = 3) const uint N_TOKS_PER_GROUP = 1;
-layout (constant_id = 4) const uint WARPSIZE = 16;
-layout (constant_id = 5) const uint MATP_PER_LOOP = 64;
-layout (constant_id = 6) const uint MATP_REDUCE = 32;
-layout (constant_id = 7) const uint WARP_V_DIM = 16;
+layout (constant_id = 0) const uint WorkGroupSize = 256;
+layout (constant_id = 1) const uint gqa_ratio = 8;
+layout (constant_id = 2) const uint HSV = 128;
+layout (constant_id = 3) const uint Wr = 1;
+layout (constant_id = 4) const uint SubGroupSize = 16;
+layout (constant_id = 5) const uint p_per_loop = 64;
+layout (constant_id = 6) const uint p_reduce = 32;
+layout (constant_id = 7) const uint Bc = 16;
+layout (constant_id = 8) const uint TM = 8;
+layout (constant_id = 9) const uint TN = 16;
+layout (constant_id = 10) const uint Flags = 0;
 
-#define TN WARPSIZE
-#define TM 8
-#define TK 16
-#define MAT_O_N (WARP_V_DIM / TM)
-#define MAT_P_M (GQA_RATIO * N_TOKS_PER_GROUP)
-#define ALIGNED_P_M ((MAT_P_M + WARPSIZE - 1) / WARPSIZE)
-#define V_HEAD_GROUPS (HEAD_DIM / WARP_V_DIM)
+#define UNALIGNED_SEQ_LEN_BIT (1 << 0)
+#define SINK_ENABLE_BIT (1 << 16)
 
-#define SUBGROUP_COUNT (GROUPSIZE / WARPSIZE)
-#define SPLIT_P_GROUPS (MATP_PER_LOOP / TK)
+const uint TK = 16;
+const uint o_row = gqa_ratio * Wr;
+const uint omat_row = (o_row + TN - 1) / TN;
+const uint o_scalar_row = (o_row + SubGroupSize - 1) / SubGroupSize;
 
-#define SLM_POOL_SIZE_O (SUBGROUP_COUNT * ALIGNED_P_M * TN * MAT_O_N * TM)
+const uint omat_col = Bc / TM;
+const uint v_head_wgs = HSV / Bc;
 
-#define P_LOAD_PER_LOOP (GROUPSIZE * 4)
-#define P_LOAD_LOOPS ((MAT_P_M * MATP_PER_LOOP + P_LOAD_PER_LOOP - 1) / P_LOAD_PER_LOOP)
-#define SLM_POOL_SIZE_P (P_LOAD_LOOPS * P_LOAD_PER_LOOP)
-#define SIZE_LOCAL_MAX (MAT_P_M * MATP_PER_LOOP / MATP_REDUCE)
-#define MAX_LOAD_LOOPS ((SIZE_LOCAL_MAX + GROUPSIZE - 1) / GROUPSIZE)
-#define SLM_POOL_SIZE_LOCAL_MAX (MAX_LOAD_LOOPS * GROUPSIZE)
-#define MAX_REDUCE_COUNT ((MAT_P_M + SUBGROUP_COUNT - 1) / SUBGROUP_COUNT)
-#define GLOBAL_MAX_SIZE (MAX_REDUCE_COUNT * SUBGROUP_COUNT)
+const uint sg_count = WorkGroupSize / SubGroupSize;
+const uint split_p_sg = p_per_loop / TK;
+const uint p_loop_col = WorkGroupSize * 4;
+const uint p_loops = (o_row * p_per_loop + p_loop_col - 1) / p_loop_col;
 
-#define SLM_POOL_SIZE_SOFTMAX_SUM (SUBGROUP_COUNT * P_LOAD_LOOPS)
+const uint size_max = o_row * p_per_loop / p_reduce;
+const uint max_loops = (size_max + WorkGroupSize - 1) / WorkGroupSize;
+const uint global_max_loops = (o_row + sg_count - 1) / sg_count;
 
-#define SLM_OFFSET_P (GLOBAL_MAX_SIZE * 2 + SLM_POOL_SIZE_SOFTMAX_SUM * 2 + SLM_POOL_SIZE_LOCAL_MAX * 2 * 2)
+#define SLM_POOL_SIZE_GLOBAL_MAX ((global_max_loops * sg_count + SubGroupSize - 1) / SubGroupSize * SubGroupSize)
+#define SLM_POOL_SIZE_O (sg_count * omat_row * TN * omat_col * TM)
+#define SLM_POOL_SIZE_MAX (max_loops * WorkGroupSize)
+#define SLM_POOL_SIZE_SOFTMAX_SUM ((sg_count * p_loops + SubGroupSize - 1) / SubGroupSize * SubGroupSize)
+#define SLM_POOL_SIZE_OOB_STAGING (sg_count * TM * TK / 4)
 
-#define SLM_OFFSET_GLOBAL_MAX 0
-#define SLM_OFFSET_SOFTMAX_SUM (SLM_OFFSET_GLOBAL_MAX + GLOBAL_MAX_SIZE)
-#define SLM_OFFSET_O (SLM_OFFSET_SOFTMAX_SUM + SLM_POOL_SIZE_SOFTMAX_SUM)
-#define SLM_OFFSET_LOCAL_MAX (GLOBAL_MAX_SIZE + SLM_POOL_SIZE_SOFTMAX_SUM)
+const uint psh_size = (p_loops * p_loop_col < omat_row * TN * p_per_loop) ? (omat_row * TN * p_per_loop) : (p_loops * p_loop_col);
+const uint psh_base = SLM_POOL_SIZE_GLOBAL_MAX * 2 + SLM_POOL_SIZE_SOFTMAX_SUM * 2 + SLM_POOL_SIZE_MAX * 2 * 2;
 
-#define P_REDUCE_VEC4 (MATP_PER_LOOP / 4)
-#define MAX_PER_LOOP (MATP_PER_LOOP / MATP_REDUCE)
-#define SLM_MAX_STRIDE (MATP_REDUCE / 4)
-#define SUB_GROUPS_PER_LINE (MATP_PER_LOOP / WARPSIZE / 4)
+const uint global_maxsh_offset = 0;
+const uint lfsh_offset = global_maxsh_offset + SLM_POOL_SIZE_GLOBAL_MAX;
+const uint osh_offset = lfsh_offset + SLM_POOL_SIZE_SOFTMAX_SUM;
+const uint maxsh_offset = SLM_POOL_SIZE_GLOBAL_MAX + SLM_POOL_SIZE_SOFTMAX_SUM;
 
-shared slm_pool_block {
-  float slm_pool_o[GLOBAL_MAX_SIZE + SLM_POOL_SIZE_SOFTMAX_SUM + SLM_POOL_SIZE_O];
-} slm_pool_f32;
+const uint p_v4 = p_per_loop / 4;
+const uint max_per_loop = p_per_loop / p_reduce;
+const uint maxsh_stride = p_reduce / 4;
+const uint sg_per_line = p_per_loop / SubGroupSize / 4;
+const bool seq_boundary_check = (Flags & UNALIGNED_SEQ_LEN_BIT) != 0;
 
-shared slm_pool_alias_block {
-  float16_t slm_pool_pv[GLOBAL_MAX_SIZE * 2 + SLM_POOL_SIZE_SOFTMAX_SUM * 2 + SLM_POOL_SIZE_LOCAL_MAX * 2 * 2 + SLM_POOL_SIZE_P * 2];
-} slm_pool_alias_f16;
+shared sh_block {
+    float F32sh[SLM_POOL_SIZE_GLOBAL_MAX + SLM_POOL_SIZE_SOFTMAX_SUM + SLM_POOL_SIZE_O];
+} F32Poolsh;
+
+shared shf16_alias {
+    float16_t F16sh[SLM_POOL_SIZE_GLOBAL_MAX * 2 + SLM_POOL_SIZE_SOFTMAX_SUM * 2 + SLM_POOL_SIZE_MAX * 2 * 2 + psh_size * 2];
+} F16Poolsh;
+
+shared shf16v4_alias {
+    f16vec4 F16V4sh[(SLM_POOL_SIZE_GLOBAL_MAX + SLM_POOL_SIZE_SOFTMAX_SUM + SLM_POOL_SIZE_MAX * 2) / 2 + SLM_POOL_SIZE_OOB_STAGING];
+} F16V4Poolsh;
 
 void main() {
-  const uint lane = gl_SubgroupInvocationID;
-  const uint v = gl_WorkGroupID.y;
-  const uint d = gl_WorkGroupID.z;
-  const uint vWarpIdx = gl_WorkGroupID.x % V_HEAD_GROUPS;
-  const uint outTokIdx = gl_WorkGroupID.x / V_HEAD_GROUPS;
-  const uint localLinearId = gl_SubgroupID;
-  const uint wgLane = localLinearId * WARPSIZE + lane;
-  const uint splitIdx = localLinearId;
-  const uint maxDim = (p.kvSeqLen + MATP_REDUCE - 1) / MATP_REDUCE;
-  const uint pDim = maxDim * MATP_REDUCE;
-  const uint kvDim = p.kvStride1;
-  const uint oDim = p.qHead * HEAD_DIM;
-  const uint offsetBaseP = (d * p.activationLength * p.qHead + v * GQA_RATIO + outTokIdx * N_TOKS_PER_GROUP * p.qHead) * pDim / 4;
-  const uint offsetBaseMax = (d * p.activationLength * p.qHead + v * GQA_RATIO + outTokIdx * N_TOKS_PER_GROUP * p.qHead) * maxDim;
-  const uint offsetBaseV = (d * p.batchStrideV + v * p.kvStride2 + vWarpIdx * WARP_V_DIM + splitIdx * TK * kvDim);
-  const uint offsetSlmP = (SLM_OFFSET_P + wgLane * 4);
-  const float fp32Min = uintBitsToFloat(0xFEFFFFFF);
-  const float fp32Max = uintBitsToFloat(0x7EFFFFFF);
-  uint offsetV = offsetBaseV;
+    const uint lane = gl_SubgroupInvocationID;
+    const uint head_v_idx = gl_WorkGroupID.y;
+    const uint batch_v_ratio = p.neq3 / p.nev3;
+    const uint batch_idx = gl_WorkGroupID.z;
+    const uint batch_v_idx = batch_idx / batch_v_ratio;
+    const uint v_idx = gl_WorkGroupID.x % v_head_wgs;
+    const uint out_row_idx = gl_WorkGroupID.x / v_head_wgs;
+    const uint sg = gl_SubgroupID;
+    const uint wg_lane = sg * SubGroupSize + lane;
+    const uint max_dim = (p.nek1 + p_reduce - 1) / p_reduce;
+    const uint p_dim = max_dim * p_reduce;
+    const uint v_dim = p.nbv1;
+    const uint o_dim = p.neq2 * HSV;
+    const uint p_base = (batch_idx * p.neq1 * p.neq2 + head_v_idx * gqa_ratio + out_row_idx * Wr * p.neq2) * p_dim / 4;
+    const uint max_base = (batch_idx * p.neq1 * p.neq2 + head_v_idx * gqa_ratio + out_row_idx * Wr * p.neq2) * max_dim;
+    const uint v_base = batch_v_idx * p.nbv3 + head_v_idx * p.nbv2 + v_idx * Bc + sg * TK * v_dim;
+    const uint o_base = (batch_idx * p.nb3 + v_idx * Bc + head_v_idx * gqa_ratio * HSV + out_row_idx * o_dim * Wr) / 4;
+    const uint psh_offset = psh_base + wg_lane * 4;
+    const float NEG_FLT_MAX_OVER_2 = uintBitsToFloat(0xFEFFFFFF);
+    const float POS_FLT_MAX_OVER_2 = uintBitsToFloat(0x7EFFFFFF);
+    uint v_offset = v_base;
 
-  coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> sums[ALIGNED_P_M][MAT_O_N];
-  f16vec4 pStateTemp[P_LOAD_LOOPS];
+    coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> OMat[omat_row][omat_col];
+    f16vec4 pv4_temp[p_loops];
 
-  float fp32CompensationP[P_LOAD_LOOPS];
+    float eMf[p_loops];
 
-  uint loadRowBase[P_LOAD_LOOPS];
-  uint loadColBase[P_LOAD_LOOPS];
-  float fp32SoftMaxSum[P_LOAD_LOOPS];
-  float fp32GlobalMaxP[P_LOAD_LOOPS];
-  uint maxRowBase[MAX_LOAD_LOOPS];
-  uint maxColBase[MAX_LOAD_LOOPS];
-  uint outOffsets[ALIGNED_P_M];
-  bool outputMask[ALIGNED_P_M];
-  float fp32SinkCoeff[ALIGNED_P_M];
+    uint p_row_base[p_loops];
+    uint p_col_base[p_loops];
+    float Lf[p_loops];
+    float global_max[p_loops];
+    uint max_row_base[max_loops];
+    uint max_col_base[max_loops];
+    uint o_offset[o_scalar_row];
+    bool o_mask[o_scalar_row];
+    float sink[o_scalar_row];
 
-  [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
-    const uint flatOffset = pm * WARPSIZE + lane;
-    const uint inGroupTokIdx = flatOffset / GQA_RATIO;
-    const uint inGroupHeadIdx = flatOffset % GQA_RATIO;
-    outputMask[pm] = (N_TOKS_PER_GROUP * outTokIdx + inGroupTokIdx < p.activationLength) && (inGroupHeadIdx < GQA_RATIO) && (inGroupTokIdx < N_TOKS_PER_GROUP);
-    outOffsets[pm] = (inGroupTokIdx * oDim + inGroupHeadIdx * HEAD_DIM) / 4;
-    if ((0x1 & p.flag) != 0) {
-      fp32SinkCoeff[pm] = sink_f32[inGroupHeadIdx + v * GQA_RATIO];
-    }
-  }
-
-  [[unroll]] for (uint maxCount = 0; maxCount < MAX_REDUCE_COUNT; maxCount++) {
-    const uint flatIdx = maxCount * SUBGROUP_COUNT + localLinearId;
-    const uint rowIdx = flatIdx % GQA_RATIO;
-    const uint tokIdx = flatIdx / GQA_RATIO;
-
-    if (tokIdx < N_TOKS_PER_GROUP) {
-      float fp32MaxReduce = fp32Min;
-      const uint maxOffset = offsetBaseMax + (tokIdx * p.qHead + rowIdx) * maxDim;
-      [[unroll]] for (uint maxReduce = 0; maxReduce < (maxDim + WARPSIZE - 1) / WARPSIZE; maxReduce++) {
-        if (maxReduce * WARPSIZE + lane < maxDim) {
-          fp32MaxReduce = max(fp32MaxReduce, max_f32[maxOffset + maxReduce * WARPSIZE + lane]);
+    [[unroll]] for (uint r = 0; r < o_scalar_row; ++r) {
+        const uint flat_offset = r * SubGroupSize + lane;
+        const uint row_idx = flat_offset / gqa_ratio;
+        const uint col_idx = flat_offset % gqa_ratio;
+        o_mask[r] = (Wr * out_row_idx + row_idx < p.neq1) && (col_idx < gqa_ratio) && (row_idx < Wr);
+        o_offset[r] = (row_idx * o_dim + col_idx * HSV) / 4;
+        if ((p.flag & SINK_ENABLE_BIT) != 0) {
+            sink[r] = data_s[col_idx + head_v_idx * gqa_ratio];
         }
-      }
-      fp32MaxReduce = subgroupMax(fp32MaxReduce);
-      if (lane == 0) {
-        slm_pool_f32.slm_pool_o[SLM_OFFSET_GLOBAL_MAX + maxCount * SUBGROUP_COUNT + localLinearId] = fp32MaxReduce;
-      }
-    } else {
-      if (lane == 0) {
-        slm_pool_f32.slm_pool_o[SLM_OFFSET_GLOBAL_MAX + maxCount * SUBGROUP_COUNT + localLinearId] = fp32Max;
-      }
-    }
-  }
-
-  barrier();
-
-  [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
-    const uint flatOffset = (pLoad * GROUPSIZE + wgLane) / P_REDUCE_VEC4;
-    const uint rowIdxFlat = flatOffset % GQA_RATIO;
-    const uint tokenIdxFlat = min(flatOffset / GQA_RATIO, N_TOKS_PER_GROUP - 1);
-    loadColBase[pLoad] = (pLoad * GROUPSIZE + wgLane) % P_REDUCE_VEC4;
-    loadRowBase[pLoad] = (tokenIdxFlat * p.qHead + rowIdxFlat);
-    fp32SoftMaxSum[pLoad] = 0.0f;
-    fp32GlobalMaxP[pLoad] = slm_pool_f32.slm_pool_o[SLM_OFFSET_GLOBAL_MAX + flatOffset];
-  }
-
-  [[unroll]] for (uint maxLoad = 0; maxLoad < MAX_LOAD_LOOPS; maxLoad++) {
-    const uint flatOffset = (maxLoad * GROUPSIZE + wgLane) / MAX_PER_LOOP;
-    const uint rowIdxFlat = flatOffset % GQA_RATIO;
-    const uint tokenIdxFlat = min(flatOffset / GQA_RATIO, N_TOKS_PER_GROUP - 1);
-    maxColBase[maxLoad] = (maxLoad * GROUPSIZE + wgLane) % MAX_PER_LOOP;
-    maxRowBase[maxLoad] = (tokenIdxFlat * p.qHead + rowIdxFlat);
-  }
-
-  [[unroll]] for (uint maxLoad = 0; maxLoad < MAX_LOAD_LOOPS; maxLoad++) {
-    const uint flatMaxOffset = maxRowBase[maxLoad] * maxDim + maxColBase[maxLoad];
-    slm_pool_f32.slm_pool_o[SLM_OFFSET_LOCAL_MAX + maxLoad * GROUPSIZE + wgLane] = max_f32[offsetBaseMax + flatMaxOffset];
-    maxColBase[maxLoad] = maxColBase[maxLoad] + MATP_PER_LOOP / MATP_REDUCE;
-  }
-
-  [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
-    const uint flatOffset = loadRowBase[pLoad] * pDim / 4 + loadColBase[pLoad];
-    pStateTemp[pLoad] = pStateVec4[offsetBaseP + flatOffset];
-  }
-
-  [[unroll]] for (uint n = 0; n < ALIGNED_P_M; n++) {
-    [[unroll]] for (uint i = 0; i < MAT_O_N; i++) {
-      sums[n][i] = coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(0.0f);
-    }
-  }
-
-  barrier();
-
-  [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
-    const uint maxOffset = (pLoad * GROUPSIZE + wgLane) / SLM_MAX_STRIDE;
-    if (loadColBase[pLoad] < pDim / 4) {
-      fp32CompensationP[pLoad] = slm_pool_f32.slm_pool_o[SLM_OFFSET_LOCAL_MAX + maxOffset];
-      float pTemp[4] = float[4](pStateTemp[pLoad].x, pStateTemp[pLoad].y, pStateTemp[pLoad].z, pStateTemp[pLoad].w);
-      float compTemp = exp(fp32CompensationP[pLoad] - fp32GlobalMaxP[pLoad]);
-      [[unroll]] for (uint kk = 0; kk < 4; kk++) {
-        pTemp[kk] = pTemp[kk] * compTemp;
-        fp32SoftMaxSum[pLoad] = fp32SoftMaxSum[pLoad] + pTemp[kk];
-        slm_pool_alias_f16.slm_pool_pv[offsetSlmP + pLoad * GROUPSIZE * 4 + kk] = float16_t(pTemp[kk]);
-      }
-    } else {
-      [[unroll]] for (uint kk = 0; kk < 4; kk++) {
-        slm_pool_alias_f16.slm_pool_pv[offsetSlmP + pLoad * GROUPSIZE * 4 + kk] = float16_t(0.0f);
-      }
     }
 
-    loadColBase[pLoad] = loadColBase[pLoad] + P_REDUCE_VEC4;
-  }
-
-  const uint loopCount = (p.kvSeqLen + MATP_PER_LOOP - 1) / MATP_PER_LOOP;
-
-  for (uint loop = 0; loop < loopCount; loop++) {
-    const uint slmPingPongLoad = (loop & 0x1);
-    const uint slmPingPongStore = ((loop + 1) & 0x1);
-
-    if (loop + 1 < loopCount) {
-      [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
-        const uint flatOffset = loadRowBase[pLoad] * pDim / 4 + loadColBase[pLoad];
-        pStateTemp[pLoad] = pStateVec4[offsetBaseP + flatOffset];
-      }
-
-      [[unroll]] for (uint maxLoad = 0; maxLoad < MAX_LOAD_LOOPS; maxLoad++) {
-        const uint flatMaxOffset = maxRowBase[maxLoad] * maxDim + maxColBase[maxLoad];
-        slm_pool_f32.slm_pool_o[SLM_OFFSET_LOCAL_MAX + slmPingPongStore * SLM_POOL_SIZE_LOCAL_MAX + maxLoad * GROUPSIZE + wgLane] = max_f32[offsetBaseMax + flatMaxOffset];
-        maxColBase[maxLoad] = maxColBase[maxLoad] + MATP_PER_LOOP / MATP_REDUCE;
-      }
-    }
-
-    barrier();
-
-    {
-      const uint coopMatOffsetP = SLM_OFFSET_P + slmPingPongLoad * SLM_POOL_SIZE_P + splitIdx * TK;
-      coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> matV[MAT_O_N];
-      [[unroll]] for (uint cc = 0; cc < MAT_O_N; cc++) {
-        coopMatLoad(
-          matV[cc],
-          vState,
-          offsetV + TM * cc,
-          kvDim,
-          gl_CooperativeMatrixLayoutColumnMajor);
-      }
-      [[unroll]] for (uint mo = 0; mo < ALIGNED_P_M; mo++) {
-        coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> matP;
-        coopMatLoad(
-          matP,
-          slm_pool_alias_f16.slm_pool_pv,
-          coopMatOffsetP + mo * TN * MATP_PER_LOOP,
-          MATP_PER_LOOP,
-          gl_CooperativeMatrixLayoutColumnMajor);
-
-        [[unroll]] for (uint no = 0; no < MAT_O_N; no++) {
-          sums[mo][no] = coopMatMulAdd(matV[no], matP, sums[mo][no]);
+    if (p.nek1 == 0) {
+        if (sg == 0) {
+            [[unroll]] for (uint c = 0; c < Bc / 4; ++c) {
+                [[unroll]] for (uint r = 0; r < o_scalar_row; ++r) {
+                    if (o_mask[r]) {
+                        data_ov4[o_base + o_offset[r] + c] = vec4(0.0f, 0.0f, 0.0f, 0.0f);
+                    }
+                }
+            }
         }
-      }
+
+        return;
     }
 
-    offsetV += MATP_PER_LOOP * kvDim;
-    if (loop * MATP_PER_LOOP + splitIdx * TK >= p.kvSeqLen) {
-      offsetV = 0;
-    }
-    if (loop + 1 < loopCount) {
-      [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
-        const uint maxOffset = (pLoad * GROUPSIZE + wgLane) / SLM_MAX_STRIDE;
-        if (loadColBase[pLoad] < pDim / 4) {
-          fp32CompensationP[pLoad] = slm_pool_f32.slm_pool_o[SLM_OFFSET_LOCAL_MAX + slmPingPongStore * SLM_POOL_SIZE_LOCAL_MAX + maxOffset];
-          float pTemp[4] = float[4](pStateTemp[pLoad].x, pStateTemp[pLoad].y, pStateTemp[pLoad].z, pStateTemp[pLoad].w);
-          float compTemp = exp(fp32CompensationP[pLoad] - fp32GlobalMaxP[pLoad]);
-          [[unroll]] for (uint kk = 0; kk < 4; kk++) {
-            pTemp[kk] = pTemp[kk] * compTemp;
-            fp32SoftMaxSum[pLoad] = fp32SoftMaxSum[pLoad] + pTemp[kk];
-            slm_pool_alias_f16.slm_pool_pv[offsetSlmP + slmPingPongStore * SLM_POOL_SIZE_P + pLoad * GROUPSIZE * 4 + kk] = float16_t(pTemp[kk]);
-          }
+    [[unroll]] for (uint idx = 0; idx < global_max_loops; ++idx) {
+        const uint flat_offset = idx * sg_count + sg;
+        const uint row_idx = flat_offset % gqa_ratio;
+        const uint tok_idx = flat_offset / gqa_ratio;
+
+        if (tok_idx < Wr) {
+            float Maxf = NEG_FLT_MAX_OVER_2;
+            const uint max_offset = max_base + (tok_idx * p.neq2 + row_idx) * max_dim;
+            [[unroll]] for (uint max_reduce = 0; max_reduce < (max_dim + SubGroupSize - 1) / SubGroupSize; ++max_reduce) {
+                if (max_reduce * SubGroupSize + lane < max_dim) {
+                    Maxf = max(Maxf, data_max[max_offset + max_reduce * SubGroupSize + lane]);
+                }
+            }
+            Maxf = subgroupMax(Maxf);
+            if (lane == 0) {
+                F32Poolsh.F32sh[global_maxsh_offset + idx * sg_count + sg] = Maxf;
+            }
         } else {
-          [[unroll]] for (uint kk = 0; kk < 4; kk++) {
-            slm_pool_alias_f16.slm_pool_pv[offsetSlmP + slmPingPongStore * SLM_POOL_SIZE_P + pLoad * GROUPSIZE * 4 + kk] = float16_t(0.0f);
-          }
+            if (lane == 0) {
+                F32Poolsh.F32sh[global_maxsh_offset + idx * sg_count + sg] = POS_FLT_MAX_OVER_2;
+            }
         }
-        loadColBase[pLoad] = loadColBase[pLoad] + P_REDUCE_VEC4;
-      }
-    }
-  }
-
-  barrier();
-
-  [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
-    fp32SoftMaxSum[pLoad] = subgroupAdd(fp32SoftMaxSum[pLoad]);
-  }
-
-  [[unroll]] for (uint mo = 0; mo < ALIGNED_P_M; mo++) {
-    [[unroll]] for (uint no = 0; no < MAT_O_N; no++) {
-      coopMatStore(
-      sums[mo][no],
-      slm_pool_f32.slm_pool_o,
-      SLM_OFFSET_O + mo * TN * WARP_V_DIM + TM * no + localLinearId * ALIGNED_P_M * TN * WARP_V_DIM,
-      WARP_V_DIM,
-      gl_CooperativeMatrixLayoutColumnMajor);
-    }
-  }
-
-  [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
-    slm_pool_f32.slm_pool_o[SLM_OFFSET_SOFTMAX_SUM + pLoad * SUBGROUP_COUNT + localLinearId] = fp32SoftMaxSum[pLoad];
-  }
-
-  barrier();
-
-  if (localLinearId == 1) {
-    const uint sumBase = SLM_OFFSET_SOFTMAX_SUM + lane * SUB_GROUPS_PER_LINE;
-    float sumTemp[ALIGNED_P_M][SUB_GROUPS_PER_LINE];
-    [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
-      [[unroll]] for (uint reduce = 0; reduce < SUB_GROUPS_PER_LINE; reduce++) {
-        sumTemp[pm][reduce] = slm_pool_f32.slm_pool_o[sumBase + pm * WARPSIZE * SUB_GROUPS_PER_LINE + reduce];
-      }
     }
 
-    [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
-      [[unroll]] for (uint reduce = 1; reduce < SUB_GROUPS_PER_LINE; reduce++) {
-        sumTemp[pm][0] = sumTemp[pm][0] + sumTemp[pm][reduce];
-      }
-    }
-
-    if ((0x1 & p.flag) != 0) {
-      [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
-        float fp32GlobalMax = slm_pool_f32.slm_pool_o[SLM_OFFSET_GLOBAL_MAX + pm * WARPSIZE + lane];
-        float sinkCompensation = fp32GlobalMax - fp32SinkCoeff[pm];
-        sinkCompensation = exp(sinkCompensation);
-        float softmaxSumTemp = sumTemp[pm][0] * sinkCompensation;
-        sumTemp[pm][0] = sumTemp[pm][0] + 1.0f / sinkCompensation;
-        sumTemp[pm][0] = 1.0f / sumTemp[pm][0];
-        sinkCompensation = sinkCompensation / (1.0f + softmaxSumTemp);
-        sumTemp[pm][0] = fp32GlobalMax < fp32SinkCoeff[pm] ? sinkCompensation : sumTemp[pm][0];
-        slm_pool_f32.slm_pool_o[SLM_OFFSET_SOFTMAX_SUM + pm * WARPSIZE + lane] = sumTemp[pm][0];
-      }
-    } else {
-      [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
-        slm_pool_f32.slm_pool_o[SLM_OFFSET_SOFTMAX_SUM + pm * WARPSIZE + lane] = 1.0f / sumTemp[pm][0];
-      }
-    }
-  }
-
-  [[unroll]] for (uint reduce = 2; reduce < SPLIT_P_GROUPS; reduce = reduce << 1 ) {
-    const uint stride = (reduce >> 1) * ALIGNED_P_M * TN * MAT_O_N * TM;
-    if ((localLinearId % reduce) == 0) {
-      const uint reduceBase = localLinearId * ALIGNED_P_M * TN * MAT_O_N * TM + SLM_OFFSET_O;
-      float sumTemp0[4];
-      float sumTemp1[4];
-      const uint reduceVec4Count = ALIGNED_P_M * TN * MAT_O_N * TM / 4 / WARPSIZE;
-      [[unroll]] for (uint totalLoads = 0; totalLoads < reduceVec4Count; totalLoads++) {
-        [[unroll]] for (uint kk = 0; kk < 4; kk++) {
-          sumTemp0[kk] = slm_pool_f32.slm_pool_o[reduceBase + totalLoads * 4 * WARPSIZE + 4 * lane + kk];
-          sumTemp1[kk] = slm_pool_f32.slm_pool_o[reduceBase + stride + totalLoads * 4 * WARPSIZE + 4 * lane + kk];
-        }
-
-        [[unroll]] for (uint kk = 0; kk < 4; kk++) {
-          sumTemp0[kk] = sumTemp0[kk] + sumTemp1[kk];
-        }
-
-        [[unroll]] for (uint kk = 0; kk < 4; kk++) {
-          slm_pool_f32.slm_pool_o[reduceBase + totalLoads * 4 * WARPSIZE + 4 * lane + kk] = sumTemp0[kk];
-        }
-      }
-    }
     barrier();
-  }
 
-  if (localLinearId == 0) {
-    const uint slmBase0 = SLM_OFFSET_O + lane * WARP_V_DIM;
-    const uint slmBase1 = slmBase0 + SPLIT_P_GROUPS / 2 * ALIGNED_P_M * TN * MAT_O_N * TM;
-
-    const uint offsetOutBase = (d * p.batchStrideO + vWarpIdx * WARP_V_DIM + v * GQA_RATIO * HEAD_DIM + outTokIdx * oDim * N_TOKS_PER_GROUP) / 4;
-    float fp32SoftMaxMul[ALIGNED_P_M];
-    float fp32Output[ALIGNED_P_M][4];
-    [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
-      fp32SoftMaxMul[pm] = slm_pool_f32.slm_pool_o[SLM_OFFSET_SOFTMAX_SUM + pm * WARPSIZE + lane];
+    [[unroll]] for (uint idx = 0; idx < p_loops; ++idx) {
+        const uint flat_offset = (idx * WorkGroupSize + wg_lane) / p_v4;
+        const uint row_idx = flat_offset % gqa_ratio;
+        const uint tok_idx = min(flat_offset / gqa_ratio, Wr - 1);
+        p_col_base[idx] = (idx * WorkGroupSize + wg_lane) % p_v4;
+        p_row_base[idx] = (tok_idx * p.neq2 + row_idx);
+        Lf[idx] = 0.0f;
+        global_max[idx] = F32Poolsh.F32sh[global_maxsh_offset + flat_offset];
     }
 
-    [[unroll]] for (uint vg = 0; vg < WARP_V_DIM / 4; vg++) {
-      [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
-        [[unroll]] for (uint vc = 0; vc < 4; vc++) {
-          fp32Output[pm][vc] = slm_pool_f32.slm_pool_o[slmBase0 + pm * WARPSIZE * WARP_V_DIM + vg * 4 + vc] * fp32SoftMaxMul[pm];
-          fp32Output[pm][vc] = fp32Output[pm][vc] + slm_pool_f32.slm_pool_o[slmBase1 + pm * WARPSIZE * WARP_V_DIM + vg * 4 + vc] * fp32SoftMaxMul[pm];
+    [[unroll]] for (uint idx = 0; idx < max_loops; ++idx) {
+        const uint flat_offset = (idx * WorkGroupSize + wg_lane) / max_per_loop;
+        const uint row_idx = flat_offset % gqa_ratio;
+        const uint tok_idx = min(flat_offset / gqa_ratio, Wr - 1);
+        max_col_base[idx] = (idx * WorkGroupSize + wg_lane) % max_per_loop;
+        max_row_base[idx] = (tok_idx * p.neq2 + row_idx);
+    }
+
+    [[unroll]] for (uint idx = 0; idx < max_loops; ++idx) {
+        const uint flat_offset = max_row_base[idx] * max_dim + max_col_base[idx];
+        if (seq_boundary_check) {
+            F32Poolsh.F32sh[maxsh_offset + idx * WorkGroupSize + wg_lane] =
+                (max_col_base[idx] < max_dim) ? data_max[max_base + flat_offset] : NEG_FLT_MAX_OVER_2;
+        } else {
+            F32Poolsh.F32sh[maxsh_offset + idx * WorkGroupSize + wg_lane] = data_max[max_base + flat_offset];
+        }
+        max_col_base[idx] = max_col_base[idx] + p_per_loop / p_reduce;
+    }
+
+    [[unroll]] for (uint idx = 0; idx < p_loops; ++idx) {
+        const uint flat_offset = p_row_base[idx] * p_dim / 4 + p_col_base[idx];
+        if (seq_boundary_check) {
+            pv4_temp[idx] = (p_col_base[idx] < p_dim / 4) ? data_pv4[p_base + flat_offset] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+        } else {
+            pv4_temp[idx] = data_pv4[p_base + flat_offset];
+        }
+    }
+
+    [[unroll]] for (uint r = 0; r < omat_row; ++r) {
+        [[unroll]] for (uint c = 0; c < omat_col; ++c) {
+            OMat[r][c] = coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(0.0f);
+        }
+    }
+
+    barrier();
+
+    [[unroll]] for (uint idx = 0; idx < p_loops; ++idx) {
+        const uint max_offset = (idx * WorkGroupSize + wg_lane) / maxsh_stride;
+        eMf[idx] = F32Poolsh.F32sh[maxsh_offset + max_offset];
+        float p_temp[4] = float[4](pv4_temp[idx].x, pv4_temp[idx].y, pv4_temp[idx].z, pv4_temp[idx].w);
+        float eMf_temp = exp(eMf[idx] - global_max[idx]);
+        [[unroll]] for (uint i = 0; i < 4; ++i) {
+            p_temp[i] = p_temp[i] * eMf_temp;
+            Lf[idx] = Lf[idx] + p_temp[i];
+            F16Poolsh.F16sh[psh_offset + idx * WorkGroupSize * 4 + i] = float16_t(p_temp[i]);
+        }
+        p_col_base[idx] = p_col_base[idx] + p_v4;
+    }
+
+    const uint loop_count = seq_boundary_check ? ((p.nek1 + p_per_loop - 1) / p_per_loop - 1) : (p.nek1 / p_per_loop);
+    uint sh_pingpong_load = 0;
+    uint sh_pingpong_store = 0;
+
+    for (uint loop = 0; loop < loop_count; ++loop) {
+        sh_pingpong_load = loop & 0x1;
+        sh_pingpong_store = (loop + 1) & 0x1;
+
+        if (seq_boundary_check || loop + 1 < loop_count) {
+            [[unroll]] for (uint idx = 0; idx < p_loops; ++idx) {
+                const uint flat_offset = p_row_base[idx] * p_dim / 4 + p_col_base[idx];
+                if (seq_boundary_check) {
+                    pv4_temp[idx] = (p_col_base[idx] < p_dim / 4) ? data_pv4[p_base + flat_offset] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+                } else {
+                    pv4_temp[idx] = data_pv4[p_base + flat_offset];
+                }
+            }
+
+            [[unroll]] for (uint idx = 0; idx < max_loops; ++idx) {
+                const uint flat_offset = max_row_base[idx] * max_dim + max_col_base[idx];
+                if (seq_boundary_check) {
+                    F32Poolsh.F32sh[maxsh_offset + sh_pingpong_store * SLM_POOL_SIZE_MAX + idx * WorkGroupSize + wg_lane] =
+                        (max_col_base[idx] < max_dim) ? data_max[max_base + flat_offset] : NEG_FLT_MAX_OVER_2;
+                } else {
+                    F32Poolsh.F32sh[maxsh_offset + sh_pingpong_store * SLM_POOL_SIZE_MAX + idx * WorkGroupSize + wg_lane] = data_max[max_base + flat_offset];
+                }
+                max_col_base[idx] = max_col_base[idx] + p_per_loop / p_reduce;
+            }
+        }
+
+        barrier();
+
+        {
+            const uint coord_p = psh_base + sh_pingpong_load * psh_size + sg * TK;
+            coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> VMat[omat_col];
+            [[unroll]] for (uint c = 0; c < omat_col; ++c) {
+                coopMatLoad(VMat[c], data_v, v_offset + TM * c, v_dim, gl_CooperativeMatrixLayoutColumnMajor);
+            }
+            [[unroll]] for (uint r = 0; r < omat_row; ++r) {
+                coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> PMat;
+                coopMatLoad(PMat, F16Poolsh.F16sh, coord_p + r * TN * p_per_loop, p_per_loop, gl_CooperativeMatrixLayoutColumnMajor);
+
+                [[unroll]] for (uint c = 0; c < omat_col; ++c) {
+                    OMat[r][c] = coopMatMulAdd(VMat[c], PMat, OMat[r][c]);
+                }
+            }
+        }
+
+        v_offset += p_per_loop * v_dim;
+
+        if (seq_boundary_check || loop + 1 < loop_count) {
+            [[unroll]] for (uint idx = 0; idx < p_loops; ++idx) {
+                const uint max_offset = (idx * WorkGroupSize + wg_lane) / maxsh_stride;
+                eMf[idx] = F32Poolsh.F32sh[maxsh_offset + sh_pingpong_store * SLM_POOL_SIZE_MAX + max_offset];
+                float p_temp[4] = float[4](pv4_temp[idx].x, pv4_temp[idx].y, pv4_temp[idx].z, pv4_temp[idx].w);
+                float eMf_temp = exp(eMf[idx] - global_max[idx]);
+                [[unroll]] for (uint i = 0; i < 4; ++i) {
+                    p_temp[i] = p_temp[i] * eMf_temp;
+                    Lf[idx] = Lf[idx] + p_temp[i];
+                    F16Poolsh.F16sh[psh_offset + sh_pingpong_store * psh_size + idx * WorkGroupSize * 4 + i] = float16_t(p_temp[i]);
+                }
+                p_col_base[idx] = p_col_base[idx] + p_v4;
+            }
+        }
+    }
+
+    barrier();
+
+    // last loop for oob handling
+    if (seq_boundary_check) {
+        const uint coord_p = psh_base + sh_pingpong_store * psh_size + sg * TK;
+        coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> VMat[omat_col];
+        coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> PMat[omat_row];
+
+        [[unroll]] for (uint r = 0; r < omat_row; ++r) {
+            coopMatLoad(PMat[r], F16Poolsh.F16sh, coord_p + r * TN * p_per_loop, p_per_loop, gl_CooperativeMatrixLayoutColumnMajor);
+        }
+
+        barrier();
+        const uint v_global_coord_base = sg * TK + loop_count * p_per_loop;
+        const uint v_global_offset = (v_base + loop_count * p_per_loop * p.nbv1) / 4;
+        const uint vv4sh_base = psh_base / 4 + sg * TM * TK / 4;
+        [[unroll]] for (uint c = 0; c < omat_col; ++c) {
+            [[unroll]] for (uint idx = 0; idx < TM * TK / SubGroupSize / 4; ++idx) {
+                const uint vv4_idx = idx * SubGroupSize + lane;
+                const uint row_idx = vv4_idx / (TM / 4);
+                const uint col_idx = vv4_idx % (TM / 4);
+                const uint vv4_offset = v_global_offset + row_idx * p.nbv1 / 4 + c * TM / 4 + col_idx;
+
+                f16vec4 vv4_temp = (v_global_coord_base + row_idx < p.nek1) ? data_vv4[vv4_offset] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+                F16V4Poolsh.F16V4sh[vv4sh_base + vv4_idx] = vv4_temp;
+            }
+
+            subgroupBarrier();
+            coopMatLoad(VMat[c], F16V4Poolsh.F16V4sh, vv4sh_base, TM / 4, gl_CooperativeMatrixLayoutColumnMajor);
+            subgroupBarrier();
+        }
+
+        [[unroll]] for (uint r = 0; r < omat_row; ++r) {
+            [[unroll]] for (uint c = 0; c < omat_col; ++c) {
+                OMat[r][c] = coopMatMulAdd(VMat[c], PMat[r], OMat[r][c]);
+            }
+        }
+    }
+
+    barrier();
+
+    [[unroll]] for (uint idx = 0; idx < p_loops; ++idx) {
+        Lf[idx] = subgroupAdd(Lf[idx]);
+    }
+
+    [[unroll]] for (uint r = 0; r < omat_row; ++r) {
+        [[unroll]] for (uint c = 0; c < omat_col; ++c) {
+            const uint coord = osh_offset + r * TN * Bc + TM * c + sg * omat_row * TN * Bc;
+            coopMatStore(OMat[r][c], F32Poolsh.F32sh, coord, Bc, gl_CooperativeMatrixLayoutColumnMajor);
+        }
+    }
+
+    [[unroll]] for (uint idx = 0; idx < p_loops; ++idx) {
+        if (lane == 0) {
+            F32Poolsh.F32sh[lfsh_offset + idx * sg_count + sg] = Lf[idx];
+        }
+    }
+
+    barrier();
+
+    if (sg == 1) {
+        const uint lf_base = lfsh_offset + lane * sg_per_line;
+        float Lfrcp_temp[o_scalar_row][sg_per_line];
+        [[unroll]] for (uint r = 0; r < o_scalar_row; ++r) {
+            [[unroll]] for (uint reduce = 0; reduce < sg_per_line; ++reduce) {
+                Lfrcp_temp[r][reduce] = F32Poolsh.F32sh[lf_base + r * SubGroupSize * sg_per_line + reduce];
+            }
         }
-      }
+
+        [[unroll]] for (uint pm = 0; pm < o_scalar_row; ++pm) {
+            [[unroll]] for (uint reduce = 1; reduce < sg_per_line; ++reduce) {
+                Lfrcp_temp[pm][0] = Lfrcp_temp[pm][0] + Lfrcp_temp[pm][reduce];
+            }
+        }
+
+        subgroupBarrier();
+
+        if ((p.flag & SINK_ENABLE_BIT) != 0) {
+            [[unroll]] for (uint r = 0; r < o_scalar_row; ++r) {
+                float Mf = F32Poolsh.F32sh[global_maxsh_offset + r * SubGroupSize + lane];
+                float sink_emf = Mf - sink[r];
+                sink_emf = exp(sink_emf);
+                float Lf_sink = Lfrcp_temp[r][0] * sink_emf;
+                Lfrcp_temp[r][0] = Lfrcp_temp[r][0] + 1.0f / sink_emf;
+                Lfrcp_temp[r][0] = 1.0f / Lfrcp_temp[r][0];
+                sink_emf = sink_emf / (1.0f + Lf_sink);
+                Lfrcp_temp[r][0] = (Mf < sink[r]) ? sink_emf : Lfrcp_temp[r][0];
+                F32Poolsh.F32sh[lfsh_offset + r * SubGroupSize + lane] = Lfrcp_temp[r][0];
+            }
+        } else {
+            [[unroll]] for (uint r = 0; r < o_scalar_row; ++r) {
+                F32Poolsh.F32sh[lfsh_offset + r * SubGroupSize + lane] = (Lfrcp_temp[r][0] == 0.0f) ? 0.0f : (1.0f / Lfrcp_temp[r][0]);
+            }
+        }
+    }
+
+    [[unroll]] for (uint reduce = 2; reduce < split_p_sg; reduce = reduce << 1) {
+        const uint stride = (reduce >> 1) * omat_row * TN * omat_col * TM;
+        if ((sg % reduce) == 0) {
+            const uint reduce_base = sg * omat_row * TN * omat_col * TM + osh_offset;
+            float sum_temp0[4];
+            float sum_temp1[4];
+            const uint reduce_count = omat_row * TN * omat_col * TM / 4 / SubGroupSize;
+            [[unroll]] for (uint idx = 0; idx < reduce_count; ++idx) {
+                [[unroll]] for (uint i = 0; i < 4; ++i) {
+                    sum_temp0[i] = F32Poolsh.F32sh[reduce_base + idx * 4 * SubGroupSize + 4 * lane + i];
+                    sum_temp1[i] = F32Poolsh.F32sh[reduce_base + stride + idx * 4 * SubGroupSize + 4 * lane + i];
+                }
+
+                [[unroll]] for (uint i = 0; i < 4; ++i) {
+                    sum_temp0[i] = sum_temp0[i] + sum_temp1[i];
+                }
+
+                [[unroll]] for (uint i = 0; i < 4; ++i) {
+                    F32Poolsh.F32sh[reduce_base + idx * 4 * SubGroupSize + 4 * lane + i] = sum_temp0[i];
+                }
+            }
+        }
+        barrier();
+    }
+
+    if (sg == 0) {
+        const uint sh_base0 = osh_offset + lane * Bc;
+        const uint sh_base1 = sh_base0 + split_p_sg / 2 * omat_row * TN * omat_col * TM;
+
+        float Lfrcp[o_scalar_row];
+        float Of[o_scalar_row][4];
+        [[unroll]] for (uint r = 0; r < o_scalar_row; ++r) {
+            Lfrcp[r] = F32Poolsh.F32sh[lfsh_offset + r * SubGroupSize + lane];
+        }
+
+        [[unroll]] for (uint c = 0; c < Bc / 4; ++c) {
+            [[unroll]] for (uint r = 0; r < o_scalar_row; ++r) {
+                [[unroll]] for (uint vc = 0; vc < 4; ++vc) {
+                    Of[r][vc] = F32Poolsh.F32sh[sh_base0 + r * SubGroupSize * Bc + c * 4 + vc] * Lfrcp[r];
+                    Of[r][vc] = Of[r][vc] + F32Poolsh.F32sh[sh_base1 + r * SubGroupSize * Bc + c * 4 + vc] * Lfrcp[r];
+                }
+            }
 
-      [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
-        if (outputMask[pm] == true) {
-          out_f32_vec4[offsetOutBase + outOffsets[pm] + vg] = vec4(fp32Output[pm][0], fp32Output[pm][1], fp32Output[pm][2], fp32Output[pm][3]);
+            [[unroll]] for (uint r = 0; r < o_scalar_row; ++r) {
+                if (o_mask[r]) {
+                    data_ov4[o_base + o_offset[r] + c] = vec4(Of[r][0], Of[r][1], Of[r][2], Of[r][3]);
+                }
+            }
         }
-      }
     }
-  }
 }

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

@@ -0,0 +1,817 @@
+#version 450
+
+#extension GL_EXT_control_flow_attributes : enable
+#extension GL_EXT_shader_16bit_storage : require
+#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
+#extension GL_KHR_cooperative_matrix : enable
+#extension GL_KHR_memory_scope_semantics : enable
+#extension GL_KHR_shader_subgroup_basic : enable
+#extension GL_EXT_shared_memory_block : enable
+#extension GL_KHR_shader_subgroup_shuffle : enable
+#extension GL_KHR_shader_subgroup_shuffle_relative : enable
+
+layout (local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer Q {float data_q[];};
+layout (binding = 0) readonly buffer QM2x4 {mat2x4 data_qm2x4[];};
+layout (binding = 1) readonly buffer K {float16_t data_k[];};
+layout (binding = 1) readonly buffer KV4 {f16vec4 data_kv4[];};
+layout (binding = 1) readonly buffer KM2x4 {f16mat2x4 data_km2x4[];};
+layout (binding = 2) readonly buffer V {float16_t data_v[];};
+layout (binding = 2) readonly buffer VV4 {f16vec4 data_vv4[];};
+layout (binding = 3) readonly buffer M {float16_t data_m[];};
+layout (binding = 3) readonly buffer MV2 {f16vec2 data_mv2[];};
+layout (binding = 3) readonly buffer MV4 {f16vec4 data_mv4[];};
+layout (binding = 4) readonly buffer S {float data_s[];};
+layout (binding = 5) buffer OV4 {vec4 data_ov4[];};
+
+layout (push_constant) uniform parameter
+{
+    uint neq1;
+    uint neq2;
+    uint neq3;
+    uint nek1;
+    uint nek2;
+    uint nek3;
+    uint nev2;
+    uint nev3;
+    uint nem1;
+    uint nem2;
+    uint nem3;
+    uint nbq1;
+    uint nbq2;
+    uint nbk1;
+    uint nbk2;
+    uint nbv1;
+    uint nbv2;
+    uint nbq3;
+    uint nbk3;
+    uint nbv3;
+    uint nb3;
+    uint flag;
+    float scale;
+} p;
+
+layout (constant_id = 0) const uint WorkGroupSize = 256;
+layout (constant_id = 1) const uint Wr = 128;
+layout (constant_id = 2) const uint Bc = 64;
+layout (constant_id = 3) const uint HSK = 128;
+layout (constant_id = 4) const uint HSV = 128;
+layout (constant_id = 5) const uint qk_ratio = 8;
+layout (constant_id = 6) const uint step_p = 32;
+layout (constant_id = 7) const uint TM = 8;
+layout (constant_id = 8) const uint TN = 16;
+layout (constant_id = 9) const uint q_reg = 96;
+layout (constant_id = 10) const uint shmem_size = 64;
+layout (constant_id = 11) const uint acc_count = 1;
+layout (constant_id = 12) const uint Flags = 0;
+layout (constant_id = 13) const uint SubGroupSize = 16;
+layout (constant_id = 14) const uint gqa_fold = 1;
+
+#define UNALIGNED_SEQ_LEN_BIT (1 << 0)
+#define UNALIGNED_Q_STRIDE_BIT (1 << 1)
+#define PARTIAL_ROW_MASK_BIT (1 << 2)
+#define INTEL_GPU_BIT (1 << 3)
+#define SINK_ENABLE_BIT (1 << 16)
+
+const bool aligned_q_stride = (Flags & UNALIGNED_Q_STRIDE_BIT) == 0;
+const bool seq_boundary_check = (Flags & UNALIGNED_SEQ_LEN_BIT) != 0;
+const bool partial_row_mask = (Flags & PARTIAL_ROW_MASK_BIT) != 0;
+const uint TK = 16;
+const uint sg_count = WorkGroupSize / SubGroupSize;
+const uint matp_count = step_p / TM;
+const uint mato_count = Bc / TM;
+const uint qk_sg = Wr / TN;
+const uint folded_row_wg = Wr / gqa_fold;
+const uint loadv_sg = (qk_sg == sg_count) ? sg_count : (sg_count - qk_sg);
+const uint k_col = (HSK % 64 == 0 && HSK <= 512 && WorkGroupSize <= 256) ? 8 : 4;
+const uint k_loop_col = WorkGroupSize * k_col;
+const uint k_loops = (step_p * HSK + k_loop_col - 1) / k_loop_col;
+const bool k_boundary_check = (step_p * HSK) % k_loop_col != 0;
+const uint v_vert_groups = step_p / TK;
+const uint v_hori_groups = loadv_sg / v_vert_groups;
+const uint v_loops = (HSV / TM + v_hori_groups - 1) / v_hori_groups;
+const bool v_boundary_check = HSV % (v_hori_groups * TM) != 0;
+const uint m_vert_groups = folded_row_wg / SubGroupSize;
+const uint m_hori_groups = sg_count / m_vert_groups;
+const uint m_loops = (step_p + m_hori_groups - 1) / m_hori_groups;
+
+const uint matq_count = q_reg / TK;
+const uint q_shmem = HSK - q_reg;
+const uint matq_shmem_count = q_shmem / TK;
+const uint max_q_reg_staging = shmem_size / 2 / Wr;
+
+const bool intel_gpu = (Flags & INTEL_GPU_BIT) != 0;
+const bool k_in_shmem = HSK <= 576;
+const bool v_in_shmem = HSK <= 128 && HSV <= 128;
+const bool mask_in_shmem = folded_row_wg <= 64;
+
+const uint oobsh_staging_size = sg_count * TM * TK;
+const uint qsh_size = Wr * q_shmem;
+const uint ksh_size = k_in_shmem ? (k_loops * k_loop_col) : oobsh_staging_size;
+const uint msh_size = mask_in_shmem ? (folded_row_wg * step_p) : 0;
+const uint psh_size = Wr * step_p;
+const uint vsh_size = v_in_shmem ? (step_p * HSV) : 0;
+const uint emfsh_size = (qk_sg == sg_count) ? 0 : Wr;
+const uint pvsh_size = intel_gpu ? 0 : (acc_count * TM * TN * sg_count);
+
+const uint qsh_base = 0;
+const uint emfsh_base = (qsh_base + qsh_size) / 2;
+const uint vsh_base = (emfsh_base + emfsh_size) * 2;
+const uint psh_base = (vsh_base + vsh_size) / 2;
+const uint ksh_base = (psh_base + psh_size) * 2;
+const uint pvsh_base = ksh_base / 2;
+const uint msh_base = ksh_base + ksh_size;
+const uint vsh_oob_staging_base = ksh_base / 4;
+
+const uint sg_mat_reduce = SubGroupSize / TN;
+const uint softmax_loops = step_p / sg_mat_reduce;
+const uint acc_loops = acc_count * TM / sg_mat_reduce;
+const uint CM1_ELEMS = (TM * TN) / SubGroupSize;
+const uint tk_sg_reduce = (TK + SubGroupSize - 1) / SubGroupSize;
+const uint o_sg_cols = (sg_count + qk_sg - 1) / qk_sg;
+const bool v_row_overshoot_check = o_sg_cols * Bc > HSV;
+
+shared sh_block {
+    float F32sh[shmem_size / 4];
+} F32Poolsh;
+
+shared shf16_alias {
+    float16_t F16sh[shmem_size / 2];
+} F16Poolsh;
+
+shared shf16v2_alias {
+    f16vec2 F16V2sh[shmem_size / 4];
+} F16V2Poolsh;
+
+shared shf16v4_alias {
+    f16vec4 F16V4sh[shmem_size / 8];
+} F16V4Poolsh;
+
+shared shf16m2x4_alias {
+    f16mat2x4 M2x4sh[shmem_size / 16];
+} F16M2x4Poolsh;
+
+void main() {
+    const uint lane = gl_SubgroupInvocationID;
+    const uint lane_c = lane % TN;
+    const uint lane_r = lane / TN;
+    const uint pv_ratio = p.neq2 / p.nev2;
+    const uint head_q_idx = gl_WorkGroupID.x;
+    const uint head_k_idx = head_q_idx * gqa_fold / qk_ratio;
+    const uint head_v_idx = head_q_idx * gqa_fold / pv_ratio;
+    const uint batch_k_ratio = p.neq3 / p.nek3;
+    const uint batch_v_ratio = p.neq3 / p.nev3;
+    const uint head_mask_idx = head_q_idx % p.nem2;
+    const uint out_r = gl_WorkGroupID.y;
+    const uint batch_idx = gl_WorkGroupID.z;
+    const uint batch_k_idx = batch_idx / batch_k_ratio;
+    const uint batch_v_idx = batch_idx / batch_v_ratio;
+    const uint batch_mask_idx = batch_idx % p.nem3;
+    const uint sg = gl_SubgroupID;
+    const uint batch_mask_stride = p.nek1 * p.nem1 * p.nem2;
+    uint sg_v_r = sg / v_vert_groups;
+    const uint sg_v_c = sg % v_vert_groups;
+    if (qk_sg < sg_count && sg >= qk_sg) {
+        sg_v_r -= qk_sg / v_vert_groups;
+    }
+    const uint sg_m_c = sg % m_hori_groups;
+    const uint sg_m_r = sg / m_hori_groups;
+    const uint sg_o_c = sg / qk_sg;
+    const uint sg_o_r = sg % qk_sg;
+    const uint sg_o_r_h = sg_o_r % gqa_fold;
+    const uint sg_o_r_row = sg_o_r / gqa_fold;
+    const uint o_dim = p.neq2 * HSV;
+    uint q_offset = out_r * folded_row_wg + sg_o_r_row * TN + lane_c;
+    q_offset = min(q_offset, p.neq1 - 1);
+    if (aligned_q_stride) {
+        q_offset = (q_offset * p.nbq1 + batch_idx * p.nbq3 + (head_q_idx * gqa_fold + sg_o_r_h) * p.nbq2) / 8 + lane_r;
+    } else {
+        q_offset = (q_offset * p.nbq1 + batch_idx * p.nbq3 + (head_q_idx * gqa_fold + sg_o_r_h) * p.nbq2) + lane_r;
+    }
+    uint k_base;
+    if (k_in_shmem) {
+        k_base = (batch_k_idx * p.nbk3 + head_k_idx * p.nbk2) / k_col;
+    } else {
+        k_base = batch_k_idx * p.nbk3 + head_k_idx * p.nbk2;
+    }
+    uint k_offset[k_loops];
+    const uint v_base = batch_v_idx * p.nbv3 + head_v_idx * p.nbv2;
+    uint v_offset;
+    if (v_in_shmem) {
+        v_offset = v_base + sg_v_r * TM + sg_v_c * p.nbv1 * TK;
+    } else {
+        v_offset = v_base + sg_o_c * Bc;
+    }
+    uint m_offset;
+    bool mask_in_boundary;
+    if (mask_in_shmem) {
+        m_offset = out_r * folded_row_wg + sg_m_r * SubGroupSize + lane;
+        mask_in_boundary = m_offset < p.nem1;
+        m_offset = min(m_offset, p.nem1 - 1);
+        m_offset = m_offset * p.nek1 + batch_mask_idx * batch_mask_stride + head_mask_idx * p.nek1 * p.nem1;
+        if (seq_boundary_check) {
+            m_offset += sg_m_c;
+        } else if (m_loops % 4 == 0) {
+            m_offset = m_offset / 4 + sg_m_c;
+        } else if (m_loops % 2 == 0) {
+            m_offset = m_offset / 2 + sg_m_c;
+        } else {
+            m_offset += sg_m_c;
+        }
+    } else {
+        m_offset = out_r * folded_row_wg + sg_o_r_row * TN + lane_c;
+        mask_in_boundary = m_offset < p.nem1;
+        m_offset = min(m_offset, p.nem1 - 1);
+        m_offset = m_offset * p.nek1 + batch_mask_idx * batch_mask_stride + lane_r * softmax_loops + head_mask_idx * p.nek1 * p.nem1;
+        if (!seq_boundary_check) {
+            m_offset /= 4;
+        }
+    }
+    const bool mask_row_valid = !partial_row_mask || mask_in_boundary;
+    const uint o_offset = (batch_idx * p.nb3 + out_r * folded_row_wg * o_dim + (head_q_idx * gqa_fold + sg_o_r_h) * HSV +
+                           sg_o_r_row * TN * o_dim + sg_o_c * Bc + lane_c * o_dim) / 4 + lane_r;
+    const uint vsh_offset = sg_v_c * TK * HSV + sg_v_r * TM * TK + vsh_base;
+    const uint msh_store_offset = (sg_m_r * SubGroupSize + lane) * step_p + msh_base;
+    const uint msh_load_offset = sg_o_r_row * TN * step_p + lane_c * step_p + lane_r * softmax_loops + msh_base;
+    const uint psh_load_offset = psh_base + sg * TN * step_p + lane_c * step_p + lane_r * softmax_loops;
+    const uint psh_store_offset = (psh_base + sg * TN * step_p + lane_c * step_p) * 2 + lane_r * softmax_loops;
+    const float NEG_FLT_MAX_OVER_2 = uintBitsToFloat(0xFEFFFFFF);
+    const float16_t NEG_F16_MAX_OVER_2 = float16_t(NEG_FLT_MAX_OVER_2);
+    const float sink = ((p.flag & SINK_ENABLE_BIT) != 0) ? data_s[head_q_idx * gqa_fold + sg_o_r_h] : 0.0f;
+    coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> QMat[matq_count];
+    uint ksh_offset[k_loops];
+    bool mask_load_k[k_loops];
+
+    if (q_reg <= max_q_reg_staging) {
+        if (sg < qk_sg) {
+            if (aligned_q_stride) {
+                [[dont_unroll]] for (uint i = 0; i < q_reg / sg_mat_reduce / 8; ++i) {
+                    const uint qsh_offset = (qsh_base + sg * q_reg * TN + lane_c * q_reg) / 8 + i * sg_mat_reduce + lane_r;
+                    mat2x4 q_temp = data_qm2x4[q_offset];
+                    q_temp *= p.scale;
+                    F16M2x4Poolsh.M2x4sh[qsh_offset] = f16mat2x4(q_temp);
+                    q_offset += sg_mat_reduce;
+                }
+            } else {
+                [[dont_unroll]] for (uint i = 0; i < q_reg / sg_mat_reduce; ++i) {
+                    const uint qsh_offset = (qsh_base + sg * q_reg * TN + lane_c * q_reg) + i * sg_mat_reduce + lane_r;
+                    float q_temp = data_q[q_offset];
+                    q_temp *= p.scale;
+                    F16Poolsh.F16sh[qsh_offset] = float16_t(q_temp);
+                    q_offset += sg_mat_reduce;
+                }
+            }
+        }
+
+        subgroupBarrier();
+
+        if (sg < qk_sg) {
+            const uint coord_q = sg * TN * q_reg + qsh_base;
+            [[unroll]] for (uint i = 0; i < matq_count; ++i) {
+                coopMatLoad(QMat[i], F16Poolsh.F16sh, coord_q + TK * i, q_reg, gl_CooperativeMatrixLayoutColumnMajor);
+            }
+        }
+    } else {
+        if (sg < qk_sg) {
+            const uint coord_q = sg * TN * TK + qsh_base;
+            [[unroll]] for (uint iq = 0; iq < q_reg / TK; ++iq) {
+                if (aligned_q_stride) {
+                    [[dont_unroll]] for (uint idx = 0; idx < TK / sg_mat_reduce / 8; ++idx) {
+                        const uint qsh_offset = (qsh_base + sg * TK * TN + lane_c * TK) / 8 + idx * sg_mat_reduce + lane_r;
+                        mat2x4 q_temp = data_qm2x4[q_offset];
+                        q_temp *= p.scale;
+                        F16M2x4Poolsh.M2x4sh[qsh_offset] = f16mat2x4(q_temp);
+                        q_offset += sg_mat_reduce;
+                    }
+                } else {
+                    [[dont_unroll]] for (uint idx = 0; idx < TK / sg_mat_reduce; ++idx) {
+                        const uint qsh_offset = (qsh_base + sg * TK * TN + lane_c * TK) + idx * sg_mat_reduce + lane_r;
+                        float q_temp = data_q[q_offset];
+                        q_temp *= p.scale;
+                        F16Poolsh.F16sh[qsh_offset] = float16_t(q_temp);
+                        q_offset += sg_mat_reduce;
+                    }
+                }
+
+                subgroupBarrier();
+                coopMatLoad(QMat[iq], F16Poolsh.F16sh, coord_q, TK, gl_CooperativeMatrixLayoutColumnMajor);
+                subgroupBarrier();
+            }
+        }
+    }
+
+    float Of[Bc / sg_mat_reduce];
+    [[unroll]] for (uint i = 0; i < Bc / sg_mat_reduce; ++i) {
+        Of[i] = 0;
+    }
+
+    barrier();
+
+    if (sg < qk_sg) {
+        if (aligned_q_stride) {
+            [[dont_unroll]] for (uint i = 0; i < q_shmem / sg_mat_reduce / 8; ++i) {
+                const uint qsh_offset = (qsh_base + sg * q_shmem * TN + lane_c * q_shmem) / 8 + i * sg_mat_reduce + lane_r;
+                mat2x4 q_temp = data_qm2x4[q_offset + i * sg_mat_reduce];
+                q_temp *= p.scale;
+                F16M2x4Poolsh.M2x4sh[qsh_offset] = f16mat2x4(q_temp);
+            }
+        } else {
+            [[dont_unroll]] for (uint i = 0; i < q_shmem / sg_mat_reduce; ++i) {
+                const uint qsh_offset = (qsh_base + sg * q_shmem * TN + lane_c * q_shmem) + i * sg_mat_reduce + lane_r;
+                float q_temp = data_q[q_offset + i * sg_mat_reduce];
+                q_temp *= p.scale;
+                F16Poolsh.F16sh[qsh_offset] = float16_t(q_temp);
+            }
+        }
+    }
+
+    if (k_in_shmem) {
+        [[unroll]] for (uint idx = 0; idx < k_loops; ++idx) {
+            const uint flat_offset = idx * WorkGroupSize + sg * SubGroupSize + lane;
+            const uint row_idx = flat_offset / (HSK / k_col);
+            const uint col_idx = flat_offset % (HSK / k_col);
+            const uint sh_chunk_c = col_idx % (TK / k_col);
+            const uint sh_chunk_r = col_idx / (TK / k_col);
+            k_offset[idx] = k_base + row_idx * p.nbk1 / k_col + col_idx;
+            ksh_offset[idx] = (ksh_base + row_idx * TK + sh_chunk_r * TK * step_p) / k_col + sh_chunk_c;
+            if (k_boundary_check) {
+                mask_load_k[idx] = row_idx < step_p;
+            }
+        }
+    }
+
+    const uint loop_count = (p.nek1 + step_p - 1) / step_p;
+    const uint k_seq_boundary = ((p.nek1 - 1) * p.nbk1 + batch_k_idx * p.nbk3 + head_k_idx * p.nbk2 + HSK) / k_col - 1;
+    float Mf = NEG_FLT_MAX_OVER_2;
+    float eMf = 0.0f;
+    float Lf = 0;
+    for (uint loop = 0; loop < loop_count; ++loop) {
+        if (k_in_shmem) {
+            if (k_boundary_check) {
+                [[unroll]] for (uint idx = 0; idx < k_loops; ++idx) {
+                    if (seq_boundary_check && loop + 1 == loop_count) {
+                        k_offset[idx] = min(k_offset[idx], k_seq_boundary);
+                    }
+
+                    if (mask_load_k[idx]) {
+                        if (k_col == 4) {
+                            F16V4Poolsh.F16V4sh[ksh_offset[idx]] = data_kv4[k_offset[idx]];
+                        } else {
+                            F16M2x4Poolsh.M2x4sh[ksh_offset[idx]] = data_km2x4[k_offset[idx]];
+                        }
+                    }
+                    k_offset[idx] += step_p * p.nbk1 / k_col;
+                }
+            } else {
+                [[unroll]] for (uint idx = 0; idx < k_loops; ++idx) {
+                    if (seq_boundary_check && loop + 1 == loop_count) {
+                        k_offset[idx] = min(k_offset[idx], k_seq_boundary);
+                    }
+
+                    if (k_col == 4) {
+                        F16V4Poolsh.F16V4sh[ksh_offset[idx]] = data_kv4[k_offset[idx]];
+                    } else {
+                        F16M2x4Poolsh.M2x4sh[ksh_offset[idx]] = data_km2x4[k_offset[idx]];
+                    }
+                    k_offset[idx] += step_p * p.nbk1 / k_col;
+                }
+            }
+        }
+
+        if (mask_in_shmem) {
+            if (seq_boundary_check) {
+                [[unroll]] for (uint idx = 0; idx < m_loops; ++idx) {
+                    if (idx * m_hori_groups + loop * step_p + sg_m_c < p.nek1) {
+                        F16Poolsh.F16sh[msh_store_offset + sg_m_c + idx * m_hori_groups] = mask_row_valid ? data_m[m_offset] : float16_t(0.0f);
+                        m_offset += m_hori_groups;
+                    } else {
+                        F16Poolsh.F16sh[msh_store_offset + sg_m_c + idx * m_hori_groups] = NEG_F16_MAX_OVER_2;
+                    }
+                }
+            } else if (m_loops % 4 == 0) {
+                [[unroll]] for (uint idx = 0; idx < m_loops / 4; ++idx) {
+                    F16V4Poolsh.F16V4sh[msh_store_offset / 4 + sg_m_c + idx * m_hori_groups] =
+                        mask_row_valid ? data_mv4[m_offset] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+                    m_offset += m_hori_groups;
+                }
+            } else if (m_loops % 2 == 0) {
+                [[unroll]] for (uint idx = 0; idx < m_loops / 2; ++idx) {
+                    F16V2Poolsh.F16V2sh[msh_store_offset / 2 + sg_m_c + idx * m_hori_groups] = mask_row_valid ? data_mv2[m_offset] : f16vec2(0.0f, 0.0f);
+                    m_offset += m_hori_groups;
+                }
+            } else {
+                [[unroll]] for (uint idx = 0; idx < m_loops; ++idx) {
+                    F16Poolsh.F16sh[msh_store_offset + sg_m_c + idx * m_hori_groups] = mask_row_valid ? data_m[m_offset] : float16_t(0.0f);
+                    m_offset += m_hori_groups;
+                }
+            }
+        }
+
+        barrier();
+        if (qk_sg == sg_count || sg < qk_sg) {
+            {
+                coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> PMat[matp_count];
+                coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> KMat[matp_count];
+                [[unroll]] for (uint pc = 0; pc < matp_count; ++pc) {
+                    PMat[pc] = coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(0.0f);
+                }
+
+                [[unroll]] for (uint d0 = 0; d0 < matq_count; ++d0) {
+                    if (k_in_shmem) {
+                        [[unroll]] for (uint pc = 0; pc < matp_count; ++pc) {
+                            const uint coord = ksh_base + step_p * TK * d0 + TM * TK * pc;
+                            coopMatLoad(KMat[pc], F16Poolsh.F16sh, coord, TK, gl_CooperativeMatrixLayoutRowMajor);
+                        }
+                    } else {
+                        [[unroll]] for (uint pc = 0; pc < matp_count; ++pc) {
+                            const uint k_base_row = loop * step_p + pc * TM;
+                            if (seq_boundary_check && loop + 1 == loop_count) {
+                                const uint k_base_offset = (k_base + k_base_row * p.nbk1 + TK * d0) / 4;
+                                [[unroll]] for (uint idx = 0; idx < TM * TK / 4 / SubGroupSize; ++idx) {
+                                    const uint flat_offset = idx * SubGroupSize + lane;
+                                    const uint row_idx = flat_offset / (TK / 4);
+                                    const uint col_idx = flat_offset % (TK / 4);
+                                    const uint k_temp_offset = k_base_offset + row_idx * p.nbk1 / 4 + col_idx;
+                                    const uint ksh_temp_offset = ksh_base / 4 + sg * TM * TK / 4 + idx * SubGroupSize + lane;
+                                    f16vec4 kv4_temp = (k_base_row + row_idx) < p.nek1 ? data_kv4[k_temp_offset] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+                                    F16V4Poolsh.F16V4sh[ksh_temp_offset] = kv4_temp;
+                                }
+
+                                subgroupBarrier();
+                                coopMatLoad(KMat[pc], F16Poolsh.F16sh, ksh_base + sg * TM * TK, TK, gl_CooperativeMatrixLayoutRowMajor);
+                                subgroupBarrier();
+                            } else {
+                                const uint coord = k_base + k_base_row * p.nbk1 + TK * d0;
+                                coopMatLoad(KMat[pc], data_k, coord, p.nbk1, gl_CooperativeMatrixLayoutRowMajor);
+                            }
+                        }
+                    }
+
+                    [[unroll]] for (uint pc = 0; pc < matp_count; ++pc) {
+                        PMat[pc] = coopMatMulAdd(KMat[pc], QMat[d0], PMat[pc]);
+                    }
+                }
+
+                [[unroll]] for (uint d0 = 0; d0 < matq_shmem_count; ++d0) {
+                    coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> QMatsh;
+                    if (k_in_shmem) {
+                        [[unroll]] for (uint pc = 0; pc < matp_count; ++pc) {
+                            const uint coord = ksh_base + matq_count * TK * step_p + step_p * TK * d0 + TM * TK * pc;
+                            coopMatLoad(KMat[pc], F16Poolsh.F16sh, coord, TK, gl_CooperativeMatrixLayoutRowMajor);
+                        }
+                    } else {
+                        [[unroll]] for (uint pc = 0; pc < matp_count; ++pc) {
+                            const uint k_base_row = loop * step_p + pc * TM;
+                            if (seq_boundary_check && loop + 1 == loop_count) {
+                                const uint k_base_offset = (k_base + k_base_row * p.nbk1 + TK * d0 + TK * matq_count) / 4;
+                                [[unroll]] for (uint idx = 0; idx < TM * TK / 4 / SubGroupSize; ++idx) {
+                                    const uint flat_offset = idx * SubGroupSize + lane;
+                                    const uint row_idx = flat_offset / (TK / 4);
+                                    const uint col_idx = flat_offset % (TK / 4);
+                                    const uint k_temp_offset = k_base_offset + row_idx * p.nbk1 / 4 + col_idx;
+                                    const uint ksh_temp_offset = ksh_base / 4 + sg * TM * TK / 4 + idx * SubGroupSize + lane;
+                                    f16vec4 kv4_temp = (k_base_row + row_idx) < p.nek1 ? data_kv4[k_temp_offset] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+                                    F16V4Poolsh.F16V4sh[ksh_temp_offset] = kv4_temp;
+                                }
+
+                                subgroupBarrier();
+                                coopMatLoad(KMat[pc], F16Poolsh.F16sh, ksh_base + sg * TM * TK, TK, gl_CooperativeMatrixLayoutRowMajor);
+                                subgroupBarrier();
+                            } else {
+                                const uint coord = k_base + k_base_row * p.nbk1 + TK * d0 + TK * matq_count;
+                                coopMatLoad(KMat[pc], data_k, coord, p.nbk1, gl_CooperativeMatrixLayoutRowMajor);
+                            }
+                        }
+                    }
+
+                    const uint coord = qsh_base + sg * TN * q_shmem + TK * d0;
+                    coopMatLoad(QMatsh, F16Poolsh.F16sh, coord, q_shmem, gl_CooperativeMatrixLayoutColumnMajor);
+                    [[unroll]] for (uint pc = 0; pc < matp_count; ++pc) {
+                        PMat[pc] = coopMatMulAdd(KMat[pc], QMatsh, PMat[pc]);
+                    }
+                }
+
+                [[unroll]] for (uint pc = 0; pc < matp_count; ++pc) {
+                    coopMatStore(PMat[pc], F32Poolsh.F32sh, psh_base + TM * pc + sg * TN * step_p, step_p, gl_CooperativeMatrixLayoutColumnMajor);
+                }
+            }
+
+            subgroupBarrier();
+
+            float Pf[softmax_loops];
+            float Mnewf = Mf;
+
+            [[unroll]] for (uint idx = 0; idx < softmax_loops; ++idx) {
+                Pf[idx] = F32Poolsh.F32sh[psh_load_offset + idx];
+            }
+
+            if (mask_in_shmem) {
+                [[unroll]] for (uint idx = 0; idx < softmax_loops; ++idx) {
+                    Pf[idx] += float(F16Poolsh.F16sh[msh_load_offset + idx]);
+                }
+            } else {
+                if (seq_boundary_check) {
+                    [[unroll]] for (uint idx = 0; idx < softmax_loops; ++idx) {
+                        if (loop * step_p + idx + lane_r * softmax_loops < p.nek1) {
+                            Pf[idx] += mask_row_valid ? float(data_m[m_offset + idx]) : 0.0f;
+                        } else {
+                            Pf[idx] += NEG_FLT_MAX_OVER_2;
+                        }
+                    }
+                    m_offset += step_p;
+                } else {
+                    [[unroll]] for (uint idx = 0; idx < softmax_loops / 4; ++idx) {
+                        f16vec4 mv4_temp = mask_row_valid ? data_mv4[m_offset + idx] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+                        float mask_f_temp[4] = float[4](float(mv4_temp.x), float(mv4_temp.y), float(mv4_temp.z), float(mv4_temp.w));
+                        [[unroll]] for (uint vec_idx = 0; vec_idx < 4; ++vec_idx) {
+                            Pf[4 * idx + vec_idx] += mask_f_temp[vec_idx];
+                        }
+                    }
+                    m_offset += step_p / 4;
+                }
+            }
+
+            if (seq_boundary_check && loop + 1 == loop_count) {
+                [[unroll]] for (uint idx = 0; idx < softmax_loops; ++idx) {
+                    if (loop * step_p + idx + lane_r * softmax_loops < p.nek1) {
+                        Mnewf = max(Mnewf, Pf[idx]);
+                    }
+                }
+            } else {
+                [[unroll]] for (uint idx = 0; idx < softmax_loops; ++idx) {
+                    Mnewf = max(Mnewf, Pf[idx]);
+                }
+            }
+
+            [[unroll]] for (uint idx = 1; idx < sg_mat_reduce; ++idx) {
+                Mnewf = max(Mnewf, subgroupShuffleDown(Mnewf, idx * TN));
+            }
+
+            Mnewf = subgroupShuffle(Mnewf, lane & (TN - 1));
+            eMf = exp(Mf - Mnewf);
+            Lf *= eMf;
+
+            if (seq_boundary_check && loop + 1 == loop_count) {
+                [[unroll]] for (uint idx = 0; idx < softmax_loops; ++idx) {
+                    if (loop * step_p + idx + lane_r * softmax_loops < p.nek1) {
+                        Pf[idx] = exp(Pf[idx] - Mnewf);
+                    } else {
+                        Pf[idx] = 0.0f;
+                    }
+                    Lf += Pf[idx];
+                }
+            } else {
+                [[unroll]] for (uint idx = 0; idx < softmax_loops; ++idx) {
+                    Pf[idx] = exp(Pf[idx] - Mnewf);
+                    Lf += Pf[idx];
+                }
+            }
+
+            if (qk_sg != sg_count && lane_r == 0) {
+                F32Poolsh.F32sh[emfsh_base + sg * TN + lane_c] = eMf;
+            }
+
+            if (SubGroupSize > TN) {
+                subgroupBarrier();
+            }
+
+            [[unroll]] for (uint idx = 0; idx < softmax_loops; ++idx) {
+                F16Poolsh.F16sh[psh_store_offset + idx] = float16_t(Pf[idx]);
+            }
+            Mf = Mnewf;
+        }
+
+        if (v_in_shmem && (qk_sg == sg_count || sg >= sg_count - loadv_sg)) {
+            if (seq_boundary_check && loop + 1 == loop_count) {
+                const uint v_global_coord_base = loop * step_p + sg_v_c * TK + lane_c;
+                if (v_boundary_check) {
+                    [[dont_unroll]] for (uint c = 0; c < tk_sg_reduce; ++c) {
+                        const uint v_coord_temp = v_global_coord_base + c * SubGroupSize;
+                        uint v_oob_offset = (v_base + v_coord_temp * p.nbv1 + sg_v_r * TM) / 4 + lane_r;
+                        uint vsh_temp_offset = vsh_offset + c * SubGroupSize + lane_c + lane_r * 4 * TK;
+                        [[dont_unroll]] for (uint r = 0; r < v_loops - 1; ++r) {
+                            [[dont_unroll]] for (uint idx = 0; idx < TM / sg_mat_reduce / 4; ++idx) {
+                                f16vec4 vv4_temp = (v_coord_temp < p.nek1) ? data_vv4[v_oob_offset + idx * sg_mat_reduce] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+                                F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 0 * TK] = vv4_temp.x;
+                                F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 1 * TK] = vv4_temp.y;
+                                F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 2 * TK] = vv4_temp.z;
+                                F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 3 * TK] = vv4_temp.w;
+                            }
+                            v_oob_offset += TM * v_hori_groups / 4;
+                            vsh_temp_offset += TM * TK * v_hori_groups;
+                        }
+                    }
+
+                    if ((sg_v_r + (v_loops - 1) * v_hori_groups) * TM < HSV) {
+                        [[dont_unroll]] for (uint c = 0; c < tk_sg_reduce; ++c) {
+                            const uint v_coord_temp = v_global_coord_base + c * SubGroupSize;
+                            const uint v_oob_offset = (v_base + v_coord_temp * p.nbv1 + sg_v_r * TM + TM * (v_loops - 1) * v_hori_groups) / 4 + lane_r;
+                            const uint vsh_temp_offset = vsh_offset + TM * (v_loops - 1) * TK * v_hori_groups + c * SubGroupSize + lane_c + lane_r * 4 * TK;
+                            [[dont_unroll]] for (uint idx = 0; idx < TM / sg_mat_reduce / 4; ++idx) {
+                                f16vec4 vv4_temp = (v_coord_temp < p.nek1) ? data_vv4[v_oob_offset + idx * sg_mat_reduce] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+                                F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 0 * TK] = vv4_temp.x;
+                                F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 1 * TK] = vv4_temp.y;
+                                F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 2 * TK] = vv4_temp.z;
+                                F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 3 * TK] = vv4_temp.w;
+                            }
+                        }
+                    }
+                } else {
+                    [[dont_unroll]] for (uint c = 0; c < tk_sg_reduce; ++c) {
+                        const uint v_coord_temp = v_global_coord_base + c * SubGroupSize;
+                        uint v_oob_offset = (v_base + v_coord_temp * p.nbv1 + sg_v_r * TM) / 4 + lane_r;
+                        uint vsh_temp_offset = vsh_offset + c * SubGroupSize + lane_c + lane_r * 4 * TK;
+                        [[dont_unroll]] for (uint r = 0; r < v_loops; ++r) {
+                            [[dont_unroll]] for (uint idx = 0; idx < TM / sg_mat_reduce / 4; ++idx) {
+                                f16vec4 vv4_temp = (v_coord_temp < p.nek1) ? data_vv4[v_oob_offset + idx * sg_mat_reduce] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+                                F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 0 * TK] = vv4_temp.x;
+                                F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 1 * TK] = vv4_temp.y;
+                                F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 2 * TK] = vv4_temp.z;
+                                F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 3 * TK] = vv4_temp.w;
+                            }
+                            v_oob_offset += TM * v_hori_groups / 4;
+                            vsh_temp_offset += TM * TK * v_hori_groups;
+                        }
+                    }
+                }
+            } else {
+                if (v_boundary_check) {
+                    [[unroll]] for (uint r = 0; r < v_loops - 1; ++r) {
+                        coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> VMat;
+                        coopMatLoad(VMat, data_v, v_offset + TM * r * v_hori_groups, p.nbv1, gl_CooperativeMatrixLayoutColumnMajor);
+                        coopMatStore(VMat, F16Poolsh.F16sh, vsh_offset + TM * r * TK * v_hori_groups, TK, gl_CooperativeMatrixLayoutRowMajor);
+                    }
+
+                    if ((sg_v_r + (v_loops - 1) * v_hori_groups) * TM < HSV) {
+                        coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> VMat;
+                        coopMatLoad(VMat, data_v, v_offset + TM * (v_loops - 1) * v_hori_groups, p.nbv1, gl_CooperativeMatrixLayoutColumnMajor);
+                        coopMatStore(VMat, F16Poolsh.F16sh, vsh_offset + TM * (v_loops - 1) * TK * v_hori_groups, TK, gl_CooperativeMatrixLayoutRowMajor);
+                    }
+                } else {
+                    [[unroll]] for (uint r = 0; r < v_loops; ++r) {
+                        coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> VMat;
+                        coopMatLoad(VMat, data_v, v_offset + TM * r * v_hori_groups, p.nbv1, gl_CooperativeMatrixLayoutColumnMajor);
+                        coopMatStore(VMat, F16Poolsh.F16sh, vsh_offset + TM * r * TK * v_hori_groups, TK, gl_CooperativeMatrixLayoutRowMajor);
+                    }
+                }
+            }
+        }
+
+        barrier();
+        {
+            const uint coord_p = psh_base * 2 + sg_o_r * TN * step_p * 2;
+            const uint coord_v_base = vsh_base + sg_o_c * mato_count * TM * TK;
+            const uint coord_pv = pvsh_base + sg * acc_count * TN * TM;
+            const uint pvsh_offset = pvsh_base + sg * acc_count * TN * TM + lane_c * acc_count * TM + lane_r * 4;
+            if (qk_sg != sg_count) {
+                eMf = F32Poolsh.F32sh[emfsh_base + sg_o_r * TN + lane_c];
+            }
+            coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> PMat[step_p / TK];
+            [[unroll]] for (uint d0 = 0; d0 < step_p / TK; ++d0) {
+                coopMatLoad(PMat[d0], F16Poolsh.F16sh, coord_p + TK * d0, 2 * step_p, gl_CooperativeMatrixLayoutColumnMajor);
+            }
+
+            [[unroll]] for (uint pv_col = 0; pv_col < mato_count / acc_count; ++pv_col) {
+                const uint coord_v = coord_v_base + pv_col * acc_count * TM * TK;
+                coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> PVMat[acc_count];
+                coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> VMat[acc_count];
+                [[unroll]] for (uint idx = 0; idx < acc_count; ++idx) {
+                    PVMat[idx] = coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(0.0f);
+                }
+
+                [[unroll]] for (uint d0 = 0; d0 < step_p / TK; ++d0) {
+                    if (v_in_shmem) {
+                        [[unroll]] for (uint idx = 0; idx < acc_count; ++idx) {
+                            coopMatLoad(VMat[idx], F16Poolsh.F16sh, coord_v + HSV * TK * d0 + TM * TK * idx, TK, gl_CooperativeMatrixLayoutRowMajor);
+                        }
+                    } else {
+                        if (seq_boundary_check && loop + 1 == loop_count) {
+                            const uint v_global_coord_base = d0 * TK + loop * step_p + lane_c;
+                            const uint vsh_staging_store = vsh_oob_staging_base + (sg * TM * TK + lane_c * TM) / 4 + lane_r;
+                            const uint vsh_staging_load = vsh_oob_staging_base * 4 + sg * TM * TK;
+                            [[unroll]] for (uint idx = 0; idx < acc_count; ++idx) {
+                                subgroupBarrier();
+                                [[dont_unroll]] for (uint c = 0; c < tk_sg_reduce; ++c) {
+                                    const uint v_coord_temp = v_global_coord_base + c * SubGroupSize;
+                                    const uint v_oob_offset = (v_base + v_coord_temp * p.nbv1 + sg_o_c * Bc + pv_col * acc_count * TM + idx * TM) / 4 + lane_r;
+
+                                    if (v_row_overshoot_check) {
+                                        const uint v_row_temp = sg_o_c * Bc + pv_col * acc_count * TM + idx * TM;
+                                        if (v_row_temp < HSV) {
+                                            [[dont_unroll]] for (uint r = 0; r < TM / sg_mat_reduce / 4; ++r) {
+                                                f16vec4 vv4_temp = (v_coord_temp < p.nek1) ? data_vv4[v_oob_offset + sg_mat_reduce * r] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+                                                F16V4Poolsh.F16V4sh[vsh_staging_store + c * SubGroupSize * TM / 4 + sg_mat_reduce * r] = vv4_temp;
+                                            }
+                                        } else {
+                                            [[dont_unroll]] for (uint r = 0; r < TM / sg_mat_reduce / 4; ++r) {
+                                                F16V4Poolsh.F16V4sh[vsh_staging_store + c * SubGroupSize * TM / 4 + sg_mat_reduce * r] = f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+                                            }
+                                        }
+                                    } else {
+                                        [[dont_unroll]] for (uint r = 0; r < TM / sg_mat_reduce / 4; ++r) {
+                                            f16vec4 vv4_temp = (v_coord_temp < p.nek1) ? data_vv4[v_oob_offset + sg_mat_reduce * r] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+                                            F16V4Poolsh.F16V4sh[vsh_staging_store + c * SubGroupSize * TM / 4 + sg_mat_reduce * r] = vv4_temp;
+                                        }
+                                    }
+                                }
+                                subgroupBarrier();
+                                coopMatLoad(VMat[idx], F16Poolsh.F16sh, vsh_staging_load, TM, gl_CooperativeMatrixLayoutColumnMajor);
+                            }
+                        } else {
+                            const uint v_coord_global = v_offset + pv_col * acc_count * TM + d0 * TK * p.nbv1;
+                            if (v_row_overshoot_check) {
+                                [[unroll]] for (uint idx = 0; idx < acc_count; ++idx) {
+                                    const uint v_row_temp = sg_o_c * Bc + pv_col * acc_count * TM + idx * TM;
+                                    if (v_row_temp < HSV) {
+                                        coopMatLoad(VMat[idx], data_v, v_coord_global + TM * idx, p.nbv1, gl_CooperativeMatrixLayoutColumnMajor);
+                                    } else {
+                                        VMat[idx] = coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA>(0.0f);
+                                    }
+                                }
+                            } else {
+                                [[unroll]] for (uint idx = 0; idx < acc_count; ++idx) {
+                                    coopMatLoad(VMat[idx], data_v, v_coord_global + TM * idx, p.nbv1, gl_CooperativeMatrixLayoutColumnMajor);
+                                }
+                            }
+                        }
+                    }
+
+                    [[unroll]] for (uint idx = 0; idx < acc_count; ++idx) {
+                        PVMat[idx] = coopMatMulAdd(VMat[idx], PMat[d0], PVMat[idx]);
+                    }
+                }
+
+                if (intel_gpu) {
+                    [[unroll]] for (uint idx = 0; idx < acc_count; ++idx) {
+                        [[unroll]] for (uint e = 0; e < CM1_ELEMS; ++e) {
+                            Of[pv_col * acc_loops + idx * CM1_ELEMS + e] = fma(Of[pv_col * acc_loops + idx * CM1_ELEMS + e],
+                                                                               eMf,
+                                                                               PVMat[idx][e]);
+                        }
+                    }
+                } else {
+                    barrier();
+
+                    [[unroll]] for (uint idx = 0; idx < acc_count; ++idx) {
+                        coopMatStore(PVMat[idx], F32Poolsh.F32sh, coord_pv + TM * idx, acc_count * TM, gl_CooperativeMatrixLayoutColumnMajor);
+                    }
+
+                    subgroupBarrier();
+
+                    [[unroll]] for (uint idx = 0; idx < acc_loops / 4; ++idx) {
+                        [[unroll]] for (uint vec_idx = 0; vec_idx < 4; ++vec_idx) {
+                            Of[pv_col * acc_loops + idx * 4 + vec_idx] = fma(Of[pv_col * acc_loops + idx * 4 + vec_idx],
+                                                                             eMf,
+                                                                             F32Poolsh.F32sh[pvsh_offset + idx * sg_mat_reduce * 4 + vec_idx]);
+                        }
+                    }
+                    barrier();
+                }
+            }
+        }
+
+        v_offset += p.nbv1 * step_p;
+        barrier();
+    }
+
+    if (sg < qk_sg) {
+        [[unroll]] for (uint idx = 1; idx < sg_mat_reduce; ++idx) {
+            Lf += subgroupShuffleDown(Lf, idx * TN);
+        }
+        if (lane_r == 0) {
+            F32Poolsh.F32sh[sg * TN + lane_c] = Mf;
+            F32Poolsh.F32sh[qk_sg * TN + sg * TN + lane_c] = Lf;
+        }
+    }
+
+    barrier();
+
+    Mf = F32Poolsh.F32sh[sg_o_r * TN + lane_c];
+    Lf = F32Poolsh.F32sh[qk_sg * TN + sg_o_r * TN + lane_c];
+
+    float Lfrcp;
+    if ((p.flag & SINK_ENABLE_BIT) != 0) {
+        float sink_emf = Mf - sink;
+        sink_emf = exp(sink_emf);
+        float Lf_sink = Lf * sink_emf;
+        Lf += 1.0f / sink_emf;
+        Lf = 1.0f / Lf;
+        sink_emf /= (1.0f + Lf_sink);
+        Lfrcp = (Mf < sink) ? sink_emf : Lf;
+    } else {
+        Lfrcp = (Lf == 0.0f) ? 0.0f : (1.0f / Lf);
+    }
+
+    [[unroll]] for (uint out_c = 0; out_c < Bc / sg_mat_reduce; ++out_c) {
+        Of[out_c] *= Lfrcp;
+    }
+
+    const uint sg_out_r = out_r * folded_row_wg + sg_o_r_row * TN + lane_c;
+    [[unroll]] for (uint o_chunk = 0; o_chunk < Bc / sg_mat_reduce / 4; ++o_chunk) {
+        const uint out_c = sg_mat_reduce * 4 * o_chunk + sg_o_c * Bc + lane_r;
+        if (sg_out_r < p.neq1 && out_c < HSV) {
+            data_ov4[o_offset + o_chunk * sg_mat_reduce] = vec4(Of[4 * o_chunk], Of[4 * o_chunk + 1], Of[4 * o_chunk + 2], Of[4 * o_chunk + 3]);
+        }
+    }
+}

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

@@ -939,6 +939,7 @@ void process_shaders() {
 #if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
     string_to_spv("fa_decode_ph1", "flash_attn_decode_phase_1.comp", {}, true, true, false, false);
     string_to_spv("fa_decode_ph2", "flash_attn_decode_phase_2.comp", {}, true, true, false, false);
+    string_to_spv("fa_prefill", "flash_attn_prefill.comp", {}, true, true, false, false);
 #endif
 
     string_to_spv("fa_sparse_compact", "flash_attn_sparse_compact.comp", {});
@@ -1074,8 +1075,6 @@ void process_shaders() {
 
     string_to_spv("silu_back_f32",  "silu_back.comp",   {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}});
 
-    string_to_spv("diag_mask_inf_f32", "diag_mask_inf.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
-
     string_to_spv("soft_max_f32", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));
     string_to_spv("soft_max_f32_f16", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}}));
     string_to_spv("soft_max_back_f32", "soft_max_back.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));

ggml/src/ggml.c

@@ -714,13 +714,13 @@ static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = {
         .from_float_ref           = (ggml_from_float_t) quantize_row_q4_1_ref,
     },
     [4] = { // GGML_TYPE_Q4_2
-        .type_name                = "DEPRECATED",
+        .type_name                = "REMOVED",
         .blck_size                = 0,
         .type_size                = 0,
         .is_quantized             = false,
     },
     [5] = { // GGML_TYPE_Q4_3
-        .type_name                = "DEPRECATED",
+        .type_name                = "REMOVED",
         .blck_size                = 0,
         .type_size                = 0,
         .is_quantized             = false,
@@ -1009,7 +1009,6 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
     "DUP",
     "ADD",
     "ADD_ID",
-    "ADD1",
     "ACC",
     "SUB",
     "MUL",
@@ -1051,8 +1050,6 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
     "GET_ROWS_BACK",
     "SET_ROWS",
     "DIAG",
-    "DIAG_MASK_INF",
-    "DIAG_MASK_ZERO",
     "SOFT_MAX",
     "SOFT_MAX_BACK",
     "ROPE",
@@ -1117,7 +1114,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
     "GLU",
 };
 
-static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102");
+static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99");
 
 static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "none",
@@ -1125,7 +1122,6 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "x",
     "x+y",
     "x[i]+y",
-    "x+y",
     "view(x,nb,offset)+=y->x",
     "x-y",
     "x*y",
@@ -1167,8 +1163,6 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "get_rows_back(x)",
     "set_rows(x)",
     "diag(x)",
-    "diag_mask_inf(x)",
-    "diag_mask_zero(x)",
     "soft_max(x)",
     "soft_max_back(x)",
     "rope(x)",
@@ -1233,7 +1227,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "glu(x)",
 };
 
-static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102");
+static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99");
 
 static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");
 
@@ -1360,12 +1354,6 @@ size_t ggml_row_size(enum ggml_type type, int64_t ne) {
     return ggml_type_size(type)*ne/ggml_blck_size(type);
 }
 
-double ggml_type_sizef(enum ggml_type type) {
-    assert(type >= 0);
-    assert(type < GGML_TYPE_COUNT);
-    return ((double)(type_traits[type].type_size))/type_traits[type].blck_size;
-}
-
 const char * ggml_type_name(enum ggml_type type) {
     assert(type >= 0);
     assert(type < GGML_TYPE_COUNT);
@@ -2172,39 +2160,6 @@ struct ggml_tensor * ggml_add_id(
     return result;
 }
 
-// ggml_add1
-
-static struct ggml_tensor * ggml_add1_impl(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        struct ggml_tensor  * b,
-        bool                  inplace) {
-    GGML_ASSERT(ggml_is_scalar(b));
-    GGML_ASSERT(ggml_is_padded_1d(a));
-
-    struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a);
-
-    result->op     = GGML_OP_ADD1;
-    result->src[0] = a;
-    result->src[1] = b;
-
-    return result;
-}
-
-struct ggml_tensor * ggml_add1(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        struct ggml_tensor  * b) {
-    return ggml_add1_impl(ctx, a, b, false);
-}
-
-struct ggml_tensor * ggml_add1_inplace(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        struct ggml_tensor  * b) {
-    return ggml_add1_impl(ctx, a, b, true);
-}
-
 // ggml_acc
 
 static struct ggml_tensor * ggml_acc_impl(
@@ -3419,16 +3374,6 @@ struct ggml_tensor * ggml_mul_mat(
     return result;
 }
 
-void ggml_mul_mat_set_prec(
-        struct ggml_tensor * a,
-        enum ggml_prec       prec) {
-    GGML_ASSERT(a->op == GGML_OP_MUL_MAT);
-
-    const int32_t prec_i32 = (int32_t) prec;
-
-    ggml_set_op_params_i32(a, 0, prec_i32);
-}
-
 void ggml_mul_mat_set_hint(
         struct ggml_tensor * a,
         enum ggml_op_hint    hint) {
@@ -4102,70 +4047,6 @@ struct ggml_tensor * ggml_diag(
     return result;
 }
 
-// ggml_diag_mask_inf
-
-static struct ggml_tensor * ggml_diag_mask_inf_impl(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        int                   n_past,
-        bool                  inplace) {
-    struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a);
-
-    int32_t params[] = { n_past };
-    ggml_set_op_params(result, params, sizeof(params));
-
-    result->op     = GGML_OP_DIAG_MASK_INF;
-    result->src[0] = a;
-
-    return result;
-}
-
-struct ggml_tensor * ggml_diag_mask_inf(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        int                   n_past) {
-    return ggml_diag_mask_inf_impl(ctx, a, n_past, false);
-}
-
-struct ggml_tensor * ggml_diag_mask_inf_inplace(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        int                   n_past) {
-    return ggml_diag_mask_inf_impl(ctx, a, n_past, true);
-}
-
-// ggml_diag_mask_zero
-
-static struct ggml_tensor * ggml_diag_mask_zero_impl(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        int                   n_past,
-        bool                  inplace) {
-    struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a);
-
-    int32_t params[] = { n_past };
-    ggml_set_op_params(result, params, sizeof(params));
-
-    result->op     = GGML_OP_DIAG_MASK_ZERO;
-    result->src[0] = a;
-
-    return result;
-}
-
-struct ggml_tensor * ggml_diag_mask_zero(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        int                   n_past) {
-    return ggml_diag_mask_zero_impl(ctx, a, n_past, false);
-}
-
-struct ggml_tensor * ggml_diag_mask_zero_inplace(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        int                   n_past) {
-    return ggml_diag_mask_zero_impl(ctx, a, n_past, true);
-}
-
 // ggml_clamp
 
 static struct ggml_tensor * ggml_clamp_impl(
@@ -4487,44 +4368,6 @@ struct ggml_tensor * ggml_rope_ext_inplace(
     );
 }
 
-struct ggml_tensor * ggml_rope_custom(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        struct ggml_tensor  * b,
-        int                   n_dims,
-        int                   mode,
-        int                   n_ctx_orig,
-        float                 freq_base,
-        float                 freq_scale,
-        float                 ext_factor,
-        float                 attn_factor,
-        float                 beta_fast,
-        float                 beta_slow) {
-    return ggml_rope_impl(
-        ctx, a, b, NULL, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale,
-        ext_factor, attn_factor, beta_fast, beta_slow, false
-    );
-}
-
-struct ggml_tensor * ggml_rope_custom_inplace(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        struct ggml_tensor  * b,
-        int                   n_dims,
-        int                   mode,
-        int                   n_ctx_orig,
-        float                 freq_base,
-        float                 freq_scale,
-        float                 ext_factor,
-        float                 attn_factor,
-        float                 beta_fast,
-        float                 beta_slow) {
-    return ggml_rope_impl(
-        ctx, a, b, NULL, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale,
-        ext_factor, attn_factor, beta_fast, beta_slow, true
-    );
-}
-
 // Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get
 // `corr_dim(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))`
 static float ggml_rope_yarn_corr_dim(int n_dims, int n_ctx_orig, float n_rot, float base) {
@@ -5264,17 +5107,6 @@ struct ggml_tensor * ggml_upscale(
     return ggml_interpolate_impl(ctx, a, a->ne[0] * scale_factor, a->ne[1] * scale_factor, a->ne[2], a->ne[3], mode);
 }
 
-struct ggml_tensor * ggml_upscale_ext(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * a,
-        int                   ne0,
-        int                   ne1,
-        int                   ne2,
-        int                   ne3,
-        enum ggml_scale_mode  mode) {
-    return ggml_interpolate_impl(ctx, a, ne0, ne1, ne2, ne3, mode);
-}
-
 struct ggml_tensor * ggml_interpolate(
         struct ggml_context * ctx,
         struct ggml_tensor  * a,
@@ -5654,16 +5486,6 @@ struct ggml_tensor * ggml_flash_attn_ext_rows(
 }
 
 
-void ggml_flash_attn_ext_set_prec(
-        struct ggml_tensor * a,
-        enum ggml_prec       prec) {
-    GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT);
-
-    const int32_t prec_i32 = (int32_t) prec;
-
-    ggml_set_op_params_i32(a, 3, prec_i32); // scale is on first pos, max_bias on second
-}
-
 enum ggml_prec ggml_flash_attn_ext_get_prec(
         const struct ggml_tensor * a) {
     GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT);
@@ -6873,22 +6695,6 @@ static void ggml_acc_or_set(
     ggml_build_forward_expand(cgraph, cgraph->grads[isrc]);
 }
 
-static void ggml_add1_or_set(
-        struct ggml_context * ctx,
-        struct ggml_cgraph  * cgraph,
-        size_t                isrc,
-        struct ggml_tensor  * tensor) {
-    struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc];
-    GGML_ASSERT(src);
-    if (cgraph->grads[isrc]) {
-        cgraph->grads[isrc] = ggml_add1_impl(ctx, cgraph->grads[isrc], tensor, cgraph->grad_accs[isrc]);
-    } else {
-        cgraph->grads[isrc] = ggml_repeat(ctx, tensor, src);
-    }
-    ggml_format_name(cgraph->grads[isrc], "grad for %s", src->name);
-    ggml_build_forward_expand(cgraph, cgraph->grads[isrc]);
-}
-
 static void ggml_sub_or_set(
         struct ggml_context * ctx,
         struct ggml_cgraph  * cgraph,
@@ -6943,14 +6749,6 @@ static void ggml_compute_backward(
                 ggml_add_or_set(ctx, cgraph, isrc1, tmp);
             }
         } break;
-        case GGML_OP_ADD1: {
-            if (src0_needs_grads) {
-                ggml_add_or_set(ctx, cgraph, isrc0, grad);
-            }
-            if (src1_needs_grads) {
-                ggml_add_or_set(ctx, cgraph, isrc1, ggml_mean(ctx, grad)); // TODO: should probably be sum instead of mean
-            }
-        } break;
         case GGML_OP_ACC: {
             if (src0_needs_grads) {
                 ggml_add_or_set(ctx, cgraph, isrc0, grad);
@@ -7023,7 +6821,7 @@ static void ggml_compute_backward(
         } break;
         case GGML_OP_SUM: {
             if (src0_needs_grads) {
-                ggml_add1_or_set(ctx, cgraph, isrc0, grad);
+                ggml_add_or_set(ctx, cgraph, isrc0, ggml_repeat(ctx, grad, src0));
             }
         } break;
         case GGML_OP_SUM_ROWS: {
@@ -7033,7 +6831,7 @@ static void ggml_compute_backward(
         } break;
         case GGML_OP_MEAN: {
             if (src0_needs_grads) {
-                ggml_add1_or_set(ctx, cgraph, isrc0, ggml_scale_impl(ctx, grad, 1.0f/src0->ne[0], 0.0, false));
+                ggml_add_or_set(ctx, cgraph, isrc0, ggml_repeat(ctx, ggml_scale_impl(ctx, grad, 1.0f/src0->ne[0], 0.0, false), src0));
             }
         } break;
         case GGML_OP_REPEAT: {
@@ -7223,20 +7021,6 @@ static void ggml_compute_backward(
                 // noop
             }
         } break;
-        case GGML_OP_DIAG_MASK_INF: {
-            if (src0_needs_grads) {
-                /* ggml_diag_mask_inf_impl() shouldn't be here */
-                /* ref:  https://github.com/ggml-org/llama.cpp/pull/4203#discussion_r1412377992 */
-                const int n_past = ((const int32_t *) tensor->op_params)[0];
-                ggml_add_or_set(ctx, cgraph, isrc0, ggml_diag_mask_zero_impl(ctx, grad, n_past, false));
-            }
-        } break;
-        case GGML_OP_DIAG_MASK_ZERO: {
-            if (src0_needs_grads) {
-                const int n_past = ((const int32_t *) tensor->op_params)[0];
-                ggml_add_or_set(ctx, cgraph, isrc0, ggml_diag_mask_zero_impl(ctx, grad, n_past, false));
-            }
-        } break;
         case GGML_OP_TRI: {
             if (src0_needs_grads) {
                 ggml_add_or_set(ctx, cgraph, isrc0, ggml_tri(ctx, grad, (enum ggml_tri_type) ggml_get_op_params_i32(tensor, 0)));

gguf-py/gguf/constants.py

@@ -998,8 +998,8 @@ class MODEL_TENSOR(IntEnum):
     HC_FFN_INJECT        = auto() # qwen4exp
     HC_ATTN_DOWN_INJECT  = auto() # qwen4exp, [hc_attn_down | hc_attn_inject]
     HC_FFN_DOWN_INJECT   = auto() # qwen4exp, [hc_ffn_down | hc_ffn_inject]
-    ATTN_QKV_GATE        = auto() # qwen4exp, [attn_qkv | attn_gate]
-    SSM_BETA_ALPHA_CAT   = auto() # qwen4exp, [ssm_beta | ssm_alpha], unlike the grouped qwen3next ssm_ba
+    ATTN_QKV_GATE        = auto() # qwen35, qwen4exp, [attn_qkv | attn_gate]
+    SSM_BETA_ALPHA_CAT   = auto() # qwen35, qwen4exp, [ssm_beta | ssm_alpha], unlike the grouped qwen3next ssm_ba
     PLE_KEY              = auto() # qwen4exp
     PLE_VALUE            = auto() # qwen4exp
     PLE_NORM_KEY         = auto() # qwen4exp
@@ -3409,6 +3409,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.ATTN_POST_NORM,
         MODEL_TENSOR.ATTN_GATE,
         MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_QKV_GATE,
         MODEL_TENSOR.FFN_GATE,
         MODEL_TENSOR.FFN_DOWN,
         MODEL_TENSOR.FFN_UP,
@@ -3418,6 +3419,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.SSM_NORM,
         MODEL_TENSOR.SSM_BETA,
         MODEL_TENSOR.SSM_ALPHA,
+        MODEL_TENSOR.SSM_BETA_ALPHA_CAT,
         MODEL_TENSOR.SSM_OUT,
         # NextN/MTP tensors - preserved but unused
         MODEL_TENSOR.NEXTN_EH_PROJ,

include/llama.h

@@ -270,7 +270,7 @@ extern "C" {
         llama_pos    *  pos;
         int32_t      *  n_seq_id;
         llama_seq_id ** seq_id;
-        int8_t       *  logits;   // TODO: rename this to "output"
+        int8_t       *  logits;
     } llama_batch;
 
     enum llama_model_kv_override_type {
@@ -389,7 +389,6 @@ extern "C" {
         float    yarn_beta_fast;   // YaRN low correction dim
         float    yarn_beta_slow;   // YaRN high correction dim
         uint32_t yarn_orig_ctx;    // YaRN original context size
-        float    defrag_thold;     // [DEPRECATED] defragment the KV cache if holes/size > thold, <= 0 disabled (default)
 
         ggml_backend_sched_eval_callback cb_eval;
         void * cb_eval_user_data;
@@ -516,11 +515,6 @@ extern "C" {
                                    void * set_tensor_data_ud, // userdata for function
               struct llama_model_params   params);
 
-    DEPRECATED(LLAMA_API struct llama_model * llama_load_model_from_file(
-                             const char * path_model,
-              struct llama_model_params   params),
-            "use llama_model_load_from_file instead");
-
     // Load a model from a file
     // If the file is split into multiple parts, the file name must follow this pattern: <name>-%05d-of-%05d.gguf
     // If the split file name does not follow this pattern, use llama_model_load_from_splits
@@ -546,20 +540,12 @@ extern "C" {
             const struct llama_model * model,
                         const char * path_model);
 
-    DEPRECATED(LLAMA_API void llama_free_model(struct llama_model * model),
-            "use llama_model_free instead");
-
     LLAMA_API void llama_model_free(struct llama_model * model);
 
     LLAMA_API struct llama_context * llama_init_from_model(
                      struct llama_model * model,
             struct llama_context_params   params);
 
-    DEPRECATED(LLAMA_API struct llama_context * llama_new_context_with_model(
-                     struct llama_model * model,
-            struct llama_context_params   params),
-            "use llama_init_from_model instead");
-
     // Frees all allocated memory
     LLAMA_API void llama_free(struct llama_context * ctx);
 
@@ -584,13 +570,6 @@ extern "C" {
     LLAMA_API uint32_t llama_n_seq_max  (const struct llama_context * ctx);
     LLAMA_API uint32_t llama_n_rs_seq   (const struct llama_context * ctx);
 
-    DEPRECATED(LLAMA_API int32_t llama_n_ctx_train(const struct llama_model * model), "use llama_model_n_ctx_train instead");
-    DEPRECATED(LLAMA_API int32_t llama_n_embd     (const struct llama_model * model), "use llama_model_n_embd instead");
-    DEPRECATED(LLAMA_API int32_t llama_n_layer    (const struct llama_model * model), "use llama_model_n_layer instead");
-    DEPRECATED(LLAMA_API int32_t llama_n_head     (const struct llama_model * model), "use llama_model_n_head instead");
-
-    DEPRECATED(LLAMA_API int32_t llama_n_vocab    (const struct llama_vocab * vocab), "use llama_vocab_n_tokens instead");
-
     LLAMA_API const struct llama_model * llama_get_model   (const struct llama_context * ctx);
     LLAMA_API           llama_memory_t   llama_get_memory  (const struct llama_context * ctx);
     LLAMA_API  enum llama_pooling_type   llama_pooling_type(const struct llama_context * ctx); // TODO: rename to llama_get_pooling_type
@@ -684,6 +663,19 @@ extern "C" {
             const char * fname_out,
             const llama_model_quantize_params * params);
 
+    //
+    // NextN / speculative draft model queries
+    //
+
+    // DFlash2 selector: number of candidate tokens the draft keeps per block position from its logits
+    // (the packed lattice also holds the top_k x top_k pairwise transition scores); 0 if the model has no selector head
+    LLAMA_API int32_t llama_model_dflash_selector_top_k(const struct llama_model * model);
+
+    // returns pointer to the target-model layer indices
+    LLAMA_API const int32_t * llama_model_target_layer_ids  (const struct llama_model * model);
+    // returns the number of extracted layers from target model
+    LLAMA_API uint32_t        llama_model_target_layer_ids_n(const struct llama_model * model);
+
     //
     // Adapters
     //
@@ -830,8 +822,6 @@ extern "C" {
     // (logits, embedding and memory)
     // Only use when saving the state, not when restoring it, otherwise the size may be too small.
     LLAMA_API size_t llama_state_get_size(struct llama_context * ctx);
-    LLAMA_API DEPRECATED(size_t llama_get_state_size(struct llama_context * ctx),
-        "use llama_state_get_size instead");
 
     // Copies the state to the specified destination address.
     // Destination needs to have allocated enough memory.
@@ -840,10 +830,6 @@ extern "C" {
             struct llama_context * ctx,
                          uint8_t * dst,
                           size_t   size);
-    LLAMA_API DEPRECATED(size_t llama_copy_state_data(
-            struct llama_context * ctx,
-                         uint8_t * dst),
-        "use llama_state_get_data instead");
 
     // Set the state reading from the specified address
     // Returns the number of bytes read
@@ -851,10 +837,6 @@ extern "C" {
             struct llama_context * ctx,
                    const uint8_t * src,
                           size_t   size);
-    LLAMA_API DEPRECATED(size_t llama_set_state_data(
-            struct llama_context * ctx,
-                   const uint8_t * src),
-        "use llama_state_set_data instead");
 
     // Save/load session file
     LLAMA_API bool llama_state_load_file(
@@ -863,25 +845,12 @@ extern "C" {
                      llama_token * tokens_out,
                           size_t   n_token_capacity,
                           size_t * n_token_count_out);
-    LLAMA_API DEPRECATED(bool llama_load_session_file(
-            struct llama_context * ctx,
-                      const char * path_session,
-                     llama_token * tokens_out,
-                          size_t   n_token_capacity,
-                          size_t * n_token_count_out),
-        "use llama_state_load_file instead");
 
     LLAMA_API bool llama_state_save_file(
             struct llama_context * ctx,
                       const char * path_session,
                const llama_token * tokens,
                           size_t   n_token_count);
-    LLAMA_API DEPRECATED(bool llama_save_session_file(
-            struct llama_context * ctx,
-                      const char * path_session,
-               const llama_token * tokens,
-                          size_t   n_token_count),
-        "use llama_state_save_file instead");
 
     // Get the exact size needed to copy the state of a single sequence
     LLAMA_API size_t llama_state_seq_get_size(
@@ -1065,6 +1034,15 @@ extern "C" {
                                                int32_t   idx,
                                      struct llama_embd   embd);
 
+    // Add an entry whose embedding is row `row` of the layer input device buffer of ctx_other.
+    // All entries of the batch must be of this kind. Positions are set with llama_batch_ext_set_pos(), n_pos_per_embd of them.
+    // Returns the batch index (>= 0), or the llama_batch_ext_add() error codes.
+    LLAMA_API int32_t llama_batch_ext_add_embd_row(struct llama_batch_ext * batch, llama_seq_id seq_id, int32_t row);
+
+    // The state of entry idx is row `row` of the h_nextn device buffer of the context set with llama_set_ctx_state_src().
+    // All entries of a batch have a state row, or none.
+    LLAMA_API bool llama_batch_ext_set_embd_state_row(struct llama_batch_ext * batch, int32_t idx, int32_t row);
+
     // Set if output embeddings should be available for the token at index idx in the batch
     // Note: for now, this is equivalent to setting the output logits
     LLAMA_API bool llama_batch_ext_set_output_embd(
@@ -1118,14 +1096,6 @@ extern "C" {
     // Returns whether the context is currently using causal attention
     LLAMA_API bool llama_get_causal_attn(const struct llama_context * ctx);
 
-    // Set whether the model is in warmup mode or not
-    // If true, all model tensors are activated during llama_decode() to load and cache their weights.
-    //
-    // note: using this can cause extra graph reallocations because it changes the graph topology with MoE models,
-    //       so it is generally not recommended to use in practice. will be removed in the future
-    DEPRECATED(LLAMA_API void llama_set_warmup(struct llama_context * ctx, bool warmup),
-            "user code should do warmup runs manually [TAG_LLAMA_GRAPH_NO_WARMUP]");
-
     // Set abort callback
     LLAMA_API void llama_set_abort_callback(struct llama_context * ctx, ggml_abort_callback abort_callback, void * abort_callback_data);
 
@@ -1170,6 +1140,57 @@ extern "C" {
     // otherwise: float[n_embd] (1-dimensional)
     LLAMA_API float * llama_get_embeddings_seq(struct llama_context * ctx, llama_seq_id seq_id);
 
+    //
+    // NextN / speculative decoding API
+    //
+
+    // Set whether the context outputs nextn embeddings or not
+    // If masked == true,  output the embeddings only for the tokens with batch.logits != 0
+    // If masked == false, output the embeddings for all tokens in the batch regardless of batch.logits
+    LLAMA_API void llama_set_embeddings_nextn(struct llama_context * ctx, bool value, bool masked);
+
+    // Select which appended NextN block the DECODER_MTP graph runs (offset past
+    // the trunk: il = n_layer() + offset). Used by the speculative NextN driver to
+    // chain multiple trained NextN heads. Default 0 (first head).
+    LLAMA_API void llama_set_nextn_layer_offset(struct llama_context * ctx, int32_t offset);
+
+    // Get all nextn embeddings, mirroring llama_get_embeddings()
+    LLAMA_API float * llama_get_embeddings_nextn(struct llama_context * ctx);
+
+    // Get the nextn embeddings for the ith token, mirroring llama_get_embeddings_ith()
+    LLAMA_API float * llama_get_embeddings_nextn_ith(struct llama_context * ctx, int32_t i);
+
+    // Set whether the context outputs the input embeddings of a specific layer
+    LLAMA_API void llama_set_embeddings_layer_inp(struct llama_context * ctx, uint32_t lid, bool value);
+
+    // Get the input embeddings of a specific layer, mirroring llama_get_embeddings()
+    LLAMA_API float * llama_get_embeddings_layer_inp(struct llama_context * ctx, uint32_t lid);
+
+    // Write the enabled layer inputs to a device buffer [n_layer_inp*n_embd, n_batch] instead of host memory.
+    // Row r holds the token at batch index r, layers in ascending layer id order.
+    LLAMA_API void llama_set_embeddings_layer_inp_dev(struct llama_context * ctx, bool value);
+
+    // Whether decode accepts entries added with llama_batch_ext_add_embd_row()
+    // (needs a DFlash/DSpark decoder context whose ctx_other has a layer input device buffer).
+    LLAMA_API bool llama_supports_embd_row(const struct llama_context * ctx);
+
+    // Write the unmasked h_nextn rows (see llama_set_embeddings_nextn) of each decode to a device buffer
+    // [n_embd_out, n_batch + n_seq_max] instead of host memory: row r < n_batch is the token at batch index r of the last decode,
+    // row n_batch + s is free for the caller, e.g. for seq s (see llama_set_embeddings_nextn_rows()).
+    // Returns false if the buffer cannot be used, the rows are then extracted to host memory.
+    LLAMA_API bool llama_set_embeddings_nextn_dev(struct llama_context * ctx, bool value);
+
+    // Copy rows of the h_nextn device buffer to or from host memory, ordered after the last decode and its readers
+    LLAMA_API bool llama_get_embeddings_nextn_rows(struct llama_context * ctx, int64_t row0, int64_t n_rows, float * out);
+    LLAMA_API bool llama_set_embeddings_nextn_rows(struct llama_context * ctx, int64_t row0, int64_t n_rows, const float * data);
+
+    // Read the state rows of batches (see llama_batch_ext_set_embd_state_row()) from the h_nextn device buffer of ctx_src.
+    // Returns whether decode accepts batches with state rows.
+    LLAMA_API bool llama_set_ctx_state_src(struct llama_context * ctx, struct llama_context * ctx_src);
+
+    // The paired context (llama_context_params::ctx_other), e.g. the target context of a draft context; may be nullptr
+    LLAMA_API struct llama_context * llama_get_ctx_other(struct llama_context * ctx);
+
     //
     // backend sampling API [EXPERIMENTAL]
     // note: use only if the llama_context was created with at least one llama_sampler_seq_config
@@ -1238,31 +1259,6 @@ extern "C" {
     LLAMA_API llama_token llama_vocab_fim_rep(const struct llama_vocab * vocab);
     LLAMA_API llama_token llama_vocab_fim_sep(const struct llama_vocab * vocab);
 
-    DEPRECATED(LLAMA_API const char * llama_token_get_text(const struct llama_vocab * vocab, llama_token token), "use llama_vocab_get_text instead");
-    DEPRECATED(LLAMA_API float llama_token_get_score(const struct llama_vocab * vocab, llama_token token), "use llama_vocab_get_score instead");
-    DEPRECATED(LLAMA_API enum llama_token_attr llama_token_get_attr(const struct llama_vocab * vocab, llama_token token), "use llama_vocab_get_attr instead");
-    DEPRECATED(LLAMA_API bool llama_token_is_eog(const struct llama_vocab * vocab, llama_token token), "use llama_vocab_is_eog instead");
-    DEPRECATED(LLAMA_API bool llama_token_is_control(const struct llama_vocab * vocab, llama_token token), "use llama_vocab_is_control instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_bos(const struct llama_vocab * vocab), "use llama_vocab_bos instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_eos(const struct llama_vocab * vocab), "use llama_vocab_eos instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_eot(const struct llama_vocab * vocab), "use llama_vocab_eot instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_cls(const struct llama_vocab * vocab), "use llama_vocab_cls instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_sep(const struct llama_vocab * vocab), "use llama_vocab_sep instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_nl (const struct llama_vocab * vocab), "use llama_vocab_nl instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_pad(const struct llama_vocab * vocab), "use llama_vocab_pad instead");
-    DEPRECATED(LLAMA_API bool llama_add_bos_token(const struct llama_vocab * vocab), "use llama_vocab_get_add_bos instead");
-    DEPRECATED(LLAMA_API bool llama_add_eos_token(const struct llama_vocab * vocab), "use llama_vocab_get_add_eos instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_fim_pre(const struct llama_vocab * vocab), "use llama_vocab_fim_pre instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_fim_suf(const struct llama_vocab * vocab), "use llama_vocab_fim_suf instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_fim_mid(const struct llama_vocab * vocab), "use llama_vocab_fim_mid instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_fim_pad(const struct llama_vocab * vocab), "use llama_vocab_fim_pad instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_fim_rep(const struct llama_vocab * vocab), "use llama_vocab_fim_rep instead");
-    DEPRECATED(LLAMA_API llama_token llama_token_fim_sep(const struct llama_vocab * vocab), "use llama_vocab_fim_sep instead");
-
-    // CLS is equivalent to BOS
-    DEPRECATED(LLAMA_API llama_token llama_vocab_cls(const struct llama_vocab * vocab), // classification
-            "use llama_vocab_bos instead");
-
     //
     // Tokenization
     //
@@ -1534,17 +1530,6 @@ extern "C" {
                           const char * grammar_str,
                           const char * grammar_root);
 
-    DEPRECATED(LLAMA_API struct llama_sampler * llama_sampler_init_grammar_lazy(
-            const struct llama_vocab * vocab,
-                          const char * grammar_str,
-                          const char * grammar_root,
-                         const char ** trigger_words,
-                                size_t num_trigger_words,
-                   const llama_token * trigger_tokens,
-                                size_t num_trigger_tokens),
-        "use llama_sampler_init_grammar_lazy_patterns instead");
-
-
     /// @details Lazy grammar sampler, introduced in https://github.com/ggml-org/llama.cpp/pull/9639
     /// @param trigger_patterns A list of patterns that will trigger the grammar sampler. Pattern will be matched from the start of the generation output, and grammar sampler will be fed content starting from its first match group.
     /// @param trigger_tokens A list of tokens that will trigger the grammar sampler. Grammar sampler will be fed content starting from the trigger token included.

models/templates/google-gemma-4-31B-it-interleaved.jinja

@@ -1,282 +0,0 @@
-{%- macro format_parameters(properties, required) -%}
-    {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}
-    {%- set ns = namespace(found_first=false) -%}
-    {%- for key, value in properties | dictsort -%}
-        {%- set add_comma = false -%}
-        {%- if key not in standard_keys -%}
-            {%- if ns.found_first %},{% endif -%}
-            {%- set ns.found_first = true -%}
-            {{ key }}:{
-            {%- if value['description'] -%}
-                description:<|"|>{{ value['description'] }}<|"|>
-                {%- set add_comma = true -%}
-            {%- endif -%}
-            {%- if value['nullable'] %}
-                {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
-                nullable:true
-            {%- endif -%}
-            {%- if value['type'] | upper == 'STRING' -%}
-                {%- if value['enum'] -%}
-                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
-                    enum:{{ format_argument(value['enum']) }}
-                {%- endif -%}
-            {%- elif value['type'] | upper == 'OBJECT' -%}
-                ,properties:{
-                {%- if value['properties'] is defined and value['properties'] is mapping -%}
-                    {{- format_parameters(value['properties'], value['required'] | default([])) -}}
-                {%- elif value is mapping -%}
-                    {{- format_parameters(value, value['required'] | default([])) -}}
-                {%- endif -%}
-                }
-                {%- if value['required'] -%}
-                    ,required:[
-                    {%- for item in value['required'] | default([]) -%}
-                        <|"|>{{- item -}}<|"|>
-                        {%- if not loop.last %},{% endif -%}
-                    {%- endfor -%}
-                    ]
-                {%- endif -%}
-            {%- elif value['type'] | upper == 'ARRAY' -%}
-                {%- if value['items'] is mapping and value['items'] -%}
-                    ,items:{
-                    {%- set ns_items = namespace(found_first=false) -%}
-                    {%- for item_key, item_value in value['items'] | dictsort -%}
-                        {%- if item_value is not none -%}
-                            {%- if ns_items.found_first %},{% endif -%}
-                            {%- set ns_items.found_first = true -%}
-                            {%- if item_key == 'properties' -%}
-                                properties:{
-                                {%- if item_value is mapping -%}
-                                    {{- format_parameters(item_value, value['items']['required'] | default([])) -}}
-                                {%- endif -%}
-                                }
-                            {%- elif item_key == 'required' -%}
-                                required:[
-                                {%- for req_item in item_value -%}
-                                    <|"|>{{- req_item -}}<|"|>
-                                    {%- if not loop.last %},{% endif -%}
-                                {%- endfor -%}
-                                ]
-                            {%- elif item_key == 'type' -%}
-                                {%- if item_value is string -%}
-                                    type:{{ format_argument(item_value | upper) }}
-                                {%- else -%}
-                                    type:{{ format_argument(item_value | map('upper') | list) }}
-                                {%- endif -%}
-                            {%- else -%}
-                                {{ item_key }}:{{ format_argument(item_value) }}
-                            {%- endif -%}
-                        {%- endif -%}
-                    {%- endfor -%}
-                    }
-                {%- endif -%}
-            {%- endif -%}
-            {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
-            type:<|"|>{{ value['type'] | upper }}<|"|>}
-        {%- endif -%}
-    {%- endfor -%}
-{%- endmacro -%}
-{%- macro format_function_declaration(tool_data) -%}
-    declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|>
-    {%- set params = tool_data['function']['parameters'] -%}
-    {%- if params -%}
-        ,parameters:{
-        {%- if params['properties'] -%}
-            properties:{ {{- format_parameters(params['properties'], params['required']) -}} },
-        {%- endif -%}
-        {%- if params['required'] -%}
-            required:[
-            {%- for item in params['required'] -%}
-                <|"|>{{- item -}}<|"|>
-                {{- ',' if not loop.last -}}
-            {%- endfor -%}
-            ],
-        {%- endif -%}
-        {%- if params['type'] -%}
-            type:<|"|>{{- params['type'] | upper -}}<|"|>}
-        {%- endif -%}
-    {%- endif -%}
-    {%- if 'response' in tool_data['function'] -%}
-        {%- set response_declaration = tool_data['function']['response'] -%}
-        ,response:{
-        {%- if response_declaration['description'] -%}
-            description:<|"|>{{- response_declaration['description'] -}}<|"|>,
-        {%- endif -%}
-        {%- if response_declaration['type'] | upper == 'OBJECT' -%}
-            type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>}
-        {%- endif -%}
-    {%- endif -%}
-    }
-{%- endmacro -%}
-{%- macro format_argument(argument, escape_keys=True) -%}
-    {%- if argument is string -%}
-        {{- '<|"|>' + argument + '<|"|>' -}}
-    {%- elif argument is boolean -%}
-        {{- 'true' if argument else 'false' -}}
-    {%- elif argument is mapping -%}
-        {{- '{' -}}
-        {%- set ns = namespace(found_first=false) -%}
-        {%- for key, value in argument | dictsort -%}
-            {%- if ns.found_first %},{% endif -%}
-            {%- set ns.found_first = true -%}
-            {%- if escape_keys -%}
-                {{- '<|"|>' + key + '<|"|>' -}}
-            {%- else -%}
-                {{- key -}}
-            {%- endif -%}
-            :{{- format_argument(value, escape_keys=escape_keys) -}}
-        {%- endfor -%}
-        {{- '}' -}}
-    {%- elif argument is sequence -%}
-        {{- '[' -}}
-        {%- for item in argument -%}
-            {{- format_argument(item, escape_keys=escape_keys) -}}
-            {%- if not loop.last %},{% endif -%}
-        {%- endfor -%}
-        {{- ']' -}}
-    {%- else -%}
-        {{- argument -}}
-    {%- endif -%}
-{%- endmacro -%}
-{%- macro strip_thinking(text) -%}
-    {%- set ns = namespace(result='') -%}
-    {%- for part in text.split('<channel|>') -%}
-        {%- if '<|channel>' in part -%}
-            {%- set ns.result = ns.result + part.split('<|channel>')[0] -%}
-        {%- else -%}
-            {%- set ns.result = ns.result + part -%}
-        {%- endif -%}
-    {%- endfor -%}
-    {{- ns.result | trim -}}
-{%- endmacro -%}
-
-{%- set ns = namespace(prev_message_type=None, last_user_message=-1) -%}
-{%- set loop_messages = messages -%}
-{{- bos_token -}}
-{#- Handle System/Tool Definitions Block -#}
-{%- if (enable_thinking is defined and enable_thinking) or tools or messages[0]['role'] in ['system', 'developer'] -%}
-    {{- '<|turn>system\n' -}}
-
-    {#- Inject Thinking token at the very top of the FIRST system turn -#}
-    {%- if enable_thinking is defined and enable_thinking -%}
-        {{- '<|think|>\n' -}}
-        {%- set ns.prev_message_type = 'think' -%}
-    {%- endif -%}
-
-    {%- if messages[0]['role'] in ['system', 'developer'] -%}
-        {{- messages[0]['content'] | trim -}}
-        {%- set loop_messages = messages[1:] -%}
-    {%- endif -%}
-
-    {%- if tools -%}
-        {%- for tool in tools %}
-            {{- '<|tool>' -}}
-            {{- format_function_declaration(tool) | trim -}}
-            {{- '<tool|>' -}}
-        {%- endfor %}
-        {%- set ns.prev_message_type = 'tool' -%}
-    {%- endif -%}
-
-    {{- '<turn|>\n' -}}
-{%- endif %}
-
-{#- Find last user message -#}
-{%- for message in loop_messages -%}
-    {%- if message['role'] == 'user' -%}
-        {%- set ns.last_user_message = loop.index0 -%}
-    {%- endif -%}
-{%- endfor -%}
-
-{#- Loop through messages -#}
-{%- for message in loop_messages -%}
-    {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
-        {%- if not (ns.prev_message_type == 'tool_response' and message['tool_calls']) -%}
-        {{- '<|turn>' + role + '\n' }}
-        {%- endif -%}
-
-        {%- set ns.prev_message_type = None -%}
-
-            {%- if message['tool_calls'] -%}
-                {#- Preserve reasoning between tool calls for model turns that come after the last user turn -#} 
-                {%- if message['reasoning_content'] and loop.index0 > ns.last_user_message -%}
-                  {{- '<|channel>thought\n' -}}
-                  {{- message['reasoning_content'] -}}
-                  {{- '<channel|>' -}}
-                {%- endif -%}
-                {%- for tool_call in message['tool_calls'] -%}
-                    {%- set function = tool_call['function'] -%}
-                    {{- '<|tool_call>call:' + function['name'] + '{' -}}
-                    {%- if function['arguments'] is mapping -%}
-                        {%- set ns_args = namespace(found_first=false) -%}
-                        {%- for key, value in function['arguments'] | dictsort -%}
-                            {%- if ns_args.found_first %},{% endif -%}
-                            {%- set ns_args.found_first = true -%}
-                            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
-                        {%- endfor -%}
-                    {%- elif function['arguments'] is string -%}
-                        {{- function['arguments'] -}}
-                    {%- endif -%}
-                    {{- '}<tool_call|>' -}}
-                {%- endfor -%}
-                {%- set ns.prev_message_type = 'tool_call' -%}
-            {%- endif -%}
-
-            {%- if message['tool_responses'] -%}
-                {#- Tool Response handling -#}
-                {%- for tool_response in message['tool_responses'] -%}
-                    {{- '<|tool_response>' -}}
-                    {%- if tool_response['response'] is mapping -%}
-                        {{- 'response:' + tool_response['name'] | default('unknown') + '{' -}}
-                        {%- for key, value in tool_response['response'] | dictsort -%}
-                            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
-                            {%- if not loop.last %},{% endif -%}
-                        {%- endfor -%}
-                        {{- '}' -}}
-                    {%- else -%}
-                        {{- 'response:' + tool_response['name'] | default('unknown') + '{value:' + format_argument(tool_response['response'], escape_keys=False) + '}' -}}
-                    {%- endif -%}
-                    {{- '<tool_response|>' -}}
-                {%- endfor -%}
-                {%- set ns.prev_message_type = 'tool_response' -%}
-            {%- endif -%}
-
-            {%- if message['content'] is string -%}
-                {%- if role == 'model' -%}
-                    {{- strip_thinking(message['content']) -}}
-                {%- else -%}
-                    {{- message['content'] | trim -}}
-                {%- endif -%}
-            {%- elif message['content'] is sequence -%}
-                {%- for item in message['content'] -%}
-                    {%- if item['type'] == 'text' -%}
-                        {%- if role == 'model' -%}
-                            {{- strip_thinking(item['text']) -}}
-                        {%- else -%}
-                            {{- item['text'] | trim -}}
-                        {%- endif -%}
-                    {%- elif item['type'] == 'image' -%}
-                        {{- '<|image|>' -}}
-                        {%- set ns.prev_message_type = 'image' -%}
-                    {%- elif item['type'] == 'audio' -%}
-                        {{- '<|audio|>' -}}
-                        {%- set ns.prev_message_type = 'audio' -%}
-                    {%- elif item['type'] == 'video' -%}
-                        {{- '<|video|>' -}}
-                        {%- set ns.prev_message_type = 'video' -%}
-                    {%- endif -%}
-                {%- endfor -%}
-            {%- endif -%}
-
-        {%- if not (message['tool_responses'] and not message['content']) -%}
-            {{- '<turn|>\n' -}}
-        {%- endif -%}
-{%- endfor -%}
-
-{%- if add_generation_prompt -%}
-    {%- if ns.prev_message_type != 'tool_response' -%}
-        {{- '<|turn>model\n' -}}
-    {%- endif -%}
-    {%- if not enable_thinking | default(false) -%}
-        {{- '<|channel>thought\n<channel|>' -}}
-    {%- endif -%}
-{%- endif -%}

src/llama-context.cpp

@@ -207,7 +207,6 @@ llama_context::llama_context(
     cparams.embeddings_nextn_masked = false;
     cparams.offload_kqv             = params.offload_kqv;
     cparams.no_perf                 = params.no_perf;
-    cparams.warmup                  = false;
 
     // +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);
@@ -1652,11 +1651,11 @@ bool llama_context::set_h_nextn_rows(int64_t row0, int64_t n_rows, const float *
     return true;
 }
 
-bool llama_context::set_state_src(llama_context * src) {
-    if (cparams.ctx_state_src != src) {
+bool llama_context::set_ctx_state_src(llama_context * ctx_src) {
+    if (cparams.ctx_state_src != ctx_src) {
         sched_need_reserve = true;
     }
-    cparams.ctx_state_src = src;
+    cparams.ctx_state_src = ctx_src;
     return supports_state_row();
 }
 
@@ -1699,14 +1698,6 @@ bool llama_context::get_causal_attn() const {
     return cparams.causal_attn;
 }
 
-void llama_context::set_warmup(bool value) {
-    LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value);
-
-    // no scheduler reserve needed because warmups are usually with small batches
-
-    cparams.warmup = value;
-}
-
 bool llama_context::set_sampler(llama_seq_id seq_id, llama_sampler * sampler) {
     if (!sampler && sampling.samplers.count(seq_id) == 0) {
         return true;
@@ -4444,7 +4435,6 @@ llama_context_params llama_context_default_params() {
         /*.yarn_beta_fast              =*/ -1.0f,
         /*.yarn_beta_slow              =*/ -1.0f,
         /*.yarn_orig_ctx               =*/ 0,
-        /*.defrag_thold                =*/ -1.0f,
         /*.cb_eval                     =*/ nullptr,
         /*.cb_eval_user_data           =*/ nullptr,
         /*.type_k                      =*/ GGML_TYPE_F16,
@@ -4613,13 +4603,6 @@ llama_context * llama_init_from_model(
     return nullptr;
 }
 
-// deprecated
-llama_context * llama_new_context_with_model(
-                 llama_model * model,
-        llama_context_params   params) {
-    return llama_init_from_model(model, params);
-}
-
 void llama_free(llama_context * ctx) {
     delete ctx;
 }
@@ -4695,10 +4678,6 @@ bool llama_get_causal_attn(const llama_context * ctx) {
     return ctx->get_causal_attn();
 }
 
-void llama_set_warmup(llama_context * ctx, bool warmup) {
-    ctx->set_warmup(warmup);
-}
-
 void llama_synchronize(llama_context * ctx) {
     ctx->synchronize();
 }
@@ -4757,8 +4736,8 @@ bool llama_set_embeddings_nextn_rows(llama_context * ctx, int64_t row0, int64_t
     return ctx->set_h_nextn_rows(row0, n_rows, data);
 }
 
-bool llama_set_embd_state_src(llama_context * ctx, llama_context * src) {
-    return ctx->set_state_src(src);
+bool llama_set_ctx_state_src(llama_context * ctx, llama_context * ctx_src) {
+    return ctx->set_ctx_state_src(ctx_src);
 }
 
 void llama_set_embeddings_layer_inp_dev(llama_context * ctx, bool value) {
@@ -4993,31 +4972,6 @@ bool llama_memory_can_shift(llama_memory_t mem) {
 
 // llama state API
 
-// deprecated
-size_t llama_get_state_size(llama_context * ctx) {
-    return llama_state_get_size(ctx);
-}
-
-// deprecated
-size_t llama_copy_state_data(llama_context * ctx, uint8_t * dst) {
-    return llama_state_get_data(ctx, dst, -1);
-}
-
-// deprecated
-size_t llama_set_state_data(llama_context * ctx, const uint8_t * src) {
-    return llama_state_set_data(ctx, src, -1);
-}
-
-// deprecated
-bool llama_load_session_file(llama_context * ctx, const char * path_session, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out) {
-    return llama_state_load_file(ctx, path_session, tokens_out, n_token_capacity, n_token_count_out);
-}
-
-// deprecated
-bool llama_save_session_file(llama_context * ctx, const char * path_session, const llama_token * tokens, size_t n_token_count) {
-    return llama_state_save_file(ctx, path_session, tokens, n_token_count);
-}
-
 // Returns the *actual* size of the state.
 // Intended to be used when saving to state to a buffer.
 size_t llama_state_get_size(llama_context * ctx) {

src/llama-context.h

@@ -133,12 +133,11 @@ struct llama_context {
     bool get_h_nextn_rows(int64_t row0, int64_t n_rows, float * out);
     bool set_h_nextn_rows(int64_t row0, int64_t n_rows, const float * data);
 
-    // state rows read from the h_nextn device buffer of src, see llama_set_embd_state_src()
-    bool set_state_src(llama_context * src);
+    // state rows read from the h_nextn device buffer of ctx_src, see llama_set_ctx_state_src()
+    bool set_ctx_state_src(llama_context * ctx_src);
     bool supports_state_row() const;
     void set_nextn_layer_offset(int32_t offset);
     void set_causal_attn(bool value);
-    void set_warmup(bool value);
 
     void set_adapters_lora(llama_adapter_lora ** adapters, size_t n_adapters, float * scales);
 

src/llama-cparams.h

@@ -50,7 +50,6 @@ struct llama_cparams {
     bool fused_kv_rows;      // unified cache: attend to the cells of each sequence only (ggml_flash_attn_ext_rows)
     bool auto_fkvr;
     bool no_perf;
-    bool warmup;             // TODO: remove [TAG_LLAMA_GRAPH_NO_WARMUP]
     bool op_offload;
     bool kv_unified;
     bool pipeline_parallel;
@@ -72,6 +71,6 @@ struct llama_cparams {
 
     llama_context * ctx_other;
 
-    // context whose h_nextn device buffer holds the state rows of a batch, see llama_set_embd_state_src()
+    // context whose h_nextn device buffer holds the state rows of a batch, see llama_set_ctx_state_src()
     llama_context * ctx_state_src;
 };

src/llama-ext.h

@@ -123,16 +123,6 @@ LLAMA_API std::vector<ggml_backend_buffer_t> llama_memory_get_buffers(const llam
 // Useful for testing purposes.
 LLAMA_API void llama_model_set_attn_rot(struct llama_model * model, bool value);
 
-// Set whether the context outputs nextn embeddings or not
-// If masked == true,  output the embeddings only for the tokens with batch.logits != 0
-// If masked == false, output the embeddings for all tokens in the batch regardless of batch.logits
-LLAMA_API void llama_set_embeddings_nextn(struct llama_context * ctx, bool value, bool masked);
-
-// Select which appended NextN block the DECODER_MTP graph runs (offset past
-// the trunk: il = n_layer() + offset). Used by the speculative NextN driver to
-// chain multiple trained NextN heads. Default 0 (first head).
-LLAMA_API void llama_set_nextn_layer_offset(struct llama_context * ctx, int32_t offset);
-
 // Marks the entries that a decision head (clef, vllm-decision2) reads, the default is 0
 // See https://github.com/ggml-org/llama.cpp/pull/29831 for details
 // A run of entries with the same value is one span, spans must be separated by entries with value 0
@@ -150,64 +140,10 @@ LLAMA_API bool llama_batch_ext_set_decision_order(struct llama_batch_ext * batch
 // Label of option i of a decision model ("<arch>.decision.labels"), nullptr past the last one
 LLAMA_API const char * llama_model_decision_label(const struct llama_model * model, uint32_t i);
 
-// mirrors:
-// LLAMA_API float * llama_get_embeddings(struct llama_context * ctx);
-LLAMA_API float * llama_get_embeddings_nextn(struct llama_context * ctx);
-
-// LLAMA_API float * llama_get_embeddings_ith(struct llama_context * ctx, int32_t i);
-LLAMA_API float * llama_get_embeddings_nextn_ith(struct llama_context * ctx, int32_t i);
-
-// Set whether the context outputs the input embeddings of a specific layer
-LLAMA_API void llama_set_embeddings_layer_inp(struct llama_context * ctx, uint32_t lid, bool value);
-
-// mirrors:
-// LLAMA_API float * llama_get_embeddings(struct llama_context * ctx);
-LLAMA_API float * llama_get_embeddings_layer_inp(struct llama_context * ctx, uint32_t lid);
-
-// Write the enabled layer inputs to a device buffer [n_layer_inp*n_embd, n_batch] instead of host memory.
-// Row r holds the token at batch index r, layers in ascending layer id order.
-LLAMA_API void llama_set_embeddings_layer_inp_dev(struct llama_context * ctx, bool value);
-
-// Whether decode accepts entries added with llama_batch_ext_add_embd_row()
-// (needs a DFlash/DSpark decoder context whose ctx_other has a layer input device buffer).
-LLAMA_API bool llama_supports_embd_row(const struct llama_context * ctx);
-
-// Add an entry whose embedding is row `row` of the layer input device buffer of ctx_other.
-// All entries of the batch must be of this kind. Positions are set with llama_batch_ext_set_pos(), n_pos_per_embd of them.
-// Returns the batch index (>= 0), or the llama_batch_ext_add() error codes.
-LLAMA_API int32_t llama_batch_ext_add_embd_row(struct llama_batch_ext * batch, llama_seq_id seq_id, int32_t row);
-
-LLAMA_API llama_context * llama_get_ctx_other(struct llama_context * ctx);
-
-// Write the unmasked h_nextn rows (see llama_set_embeddings_nextn) of each decode to a device buffer
-// [n_embd_out, n_batch + n_seq_max] instead of host memory: row r < n_batch is the token at batch index r of the last decode,
-// row n_batch + s is free for the caller, e.g. for seq s (see llama_set_embeddings_nextn_rows()).
-// Returns false if the buffer cannot be used, the rows are then extracted to host memory.
-LLAMA_API bool llama_set_embeddings_nextn_dev(struct llama_context * ctx, bool value);
-
-// Copy rows of the h_nextn device buffer to or from host memory, ordered after the last decode and its readers
-LLAMA_API bool llama_get_embeddings_nextn_rows(struct llama_context * ctx, int64_t row0, int64_t n_rows, float * out);
-LLAMA_API bool llama_set_embeddings_nextn_rows(struct llama_context * ctx, int64_t row0, int64_t n_rows, const float * data);
-
-// Read the state rows of batches (see llama_batch_ext_set_embd_state_row()) from the h_nextn device buffer of src.
-// Returns whether decode accepts batches with state rows.
-LLAMA_API bool llama_set_embd_state_src(struct llama_context * ctx, struct llama_context * src);
-
-// The state of entry idx is row `row` of the h_nextn device buffer of the context set with llama_set_embd_state_src().
-// All entries of a batch have a state row, or none.
-LLAMA_API bool llama_batch_ext_set_embd_state_row(struct llama_batch_ext * batch, int32_t idx, int32_t row);
-
 //
 // model/context data extraction
 //
 
-LLAMA_API int32_t llama_model_dflash_selector_top_k(const struct llama_model * model);
-
-// returns pointer to the target-model layer indices
-LLAMA_API const int32_t * llama_model_target_layer_ids  (const struct llama_model * model);
-// returns the number of extracted layers from target model
-LLAMA_API uint32_t        llama_model_target_layer_ids_n(const struct llama_model * model);
-
 // retrieves the whole token embedding matrix in F32 format (n_embd * n_vocab)
 // returns total number of elements or 0 on error
 // if out is nullptr, returns the number of tokens without writing to out

src/llama-graph.cpp

@@ -1569,7 +1569,7 @@ llm_graph_context::llm_graph_context(const llm_graph_params & params) :
     n_embd_head_v    (hparams.n_embd_head_v()),
     n_embd_v_gqa     (hparams.n_embd_v_gqa()),
     n_expert         (hparams.n_expert),
-    n_expert_used    (cparams.warmup ? hparams.n_expert : hparams.n_expert_used()),
+    n_expert_used    (hparams.n_expert_used()),
     freq_base        (cparams.rope_freq_base),
     freq_scale       (cparams.rope_freq_scale),
     ext_factor       (cparams.yarn_ext_factor),
@@ -2687,7 +2687,16 @@ llm_graph_moe_cache_slots llm_graph_context::build_moe_cache_slots(
         return { nullptr, nullptr };
     }
     for (ggml_tensor * w : { up_exps, gate_exps, down_exps, gate_up_exps }) {
-        if (w != nullptr && moe_cache->get_experts(w) == nullptr) {
+        if (w == nullptr) {
+            continue;
+        }
+        ggml_tensor * experts = moe_cache->get_experts(w);
+        if (experts == nullptr) {
+            return { nullptr, nullptr };
+        }
+        // the device of the cache must compute the matmul, else the scheduler copies the cache to another backend
+        ggml_tensor * b = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, experts->ne[0], w == down_exps ? selected_experts->ne[0] : 1, selected_experts->ne[1]);
+        if (!ggml_backend_supports_op(moe_cache->backend(il), ggml_mul_mat_id(ctx0, experts, b, selected_experts))) {
             return { nullptr, nullptr };
         }
     }

src/llama-graph.h

@@ -1267,7 +1267,7 @@ struct llm_graph_context {
     // tok_scale: applied to token rows only
     ggml_tensor * build_inp_embd(ggml_tensor * tok_embd, float tok_scale = 1.0f) const;
 
-    // the state input of an MTP graph: rows of the h_nextn device buffer of ctx_state_src (see llama_set_embd_state_src()),
+    // the state input of an MTP graph: rows of the h_nextn device buffer of ctx_state_src (see llama_set_ctx_state_src()),
     // or uploaded from ubatch.embd_state
     ggml_tensor * build_inp_embd_state(llm_graph_input_embd_h * inp, int64_t n_embd_state) const;
 

src/llama-model.cpp

@@ -424,6 +424,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
     static const std::regex pattern_kv_weight       ("blk\\.\\d*\\.attn_(k|v).weight");
     static const std::regex pattern_v_exps_weight   ("blk\\.\\d*\\.attn_v_exps.weight"); // K2 Horizon MoVA
     static const std::regex pattern_qkv_weight      ("blk\\.\\d*\\.attn_qkv.weight");
+    static const std::regex pattern_qkv_gate_weight ("blk\\.\\d*\\.attn_qkv_gate.weight");
     static const std::regex pattern_q_bias          ("blk\\.\\d*\\.attn_q\\.bias");
     static const std::regex pattern_kv_bias         ("blk\\.\\d*\\.attn_(k|v)\\.bias");
     static const std::regex pattern_qkv_bias        ("blk\\.\\d*\\.attn_qkv.bias");
@@ -446,6 +447,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
     static const std::regex pattern_ssm_alpha       ("blk\\.\\d*\\.ssm_alpha.weight");
     static const std::regex pattern_ssm_beta        ("blk\\.\\d*\\.ssm_beta.weight");
     static const std::regex pattern_ssm_beta_alpha  ("blk\\.\\d*\\.ssm_ba.weight");
+    static const std::regex pattern_ssm_beta_alpha_cat("blk\\.\\d*\\.ssm_beta_alpha.weight");
     static const std::regex pattern_r_cache         ("cache_r_l\\d*");
     static const std::regex pattern_ple_r_cache     ("cache_ple_r_l\\d*");
     static const std::regex pattern_s_cache         ("cache_s_l\\d*");
@@ -593,7 +595,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
         if (std::regex_match(tensor_name, pattern_q_bias) || std::regex_match(tensor_name, pattern_kv_bias)) {
             return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "attn_output.weight", "ssm_out.weight");
         }
-        if (std::regex_match(tensor_name, pattern_qkv_weight)) {
+        if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_gate_weight)) {
             return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight", "ssm_out.weight");
         }
         if ( std::regex_match(tensor_name, pattern_qkv_bias)) {
@@ -620,7 +622,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
             return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "ssm_out.weight");
         }
         if (std::regex_match(tensor_name, pattern_ssm_alpha) || std::regex_match(tensor_name, pattern_ssm_beta) ||
-                std::regex_match(tensor_name, pattern_ssm_beta_alpha)) {
+                std::regex_match(tensor_name, pattern_ssm_beta_alpha) || std::regex_match(tensor_name, pattern_ssm_beta_alpha_cat)) {
             return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "ssm_out.weight");
         }
         if (std::regex_match(tensor_name, pattern_r_cache) || std::regex_match(tensor_name, pattern_s_cache)) {
@@ -715,6 +717,14 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
                 }
             } else {
                 const int64_t head_ratio = n_v_heads / n_k_heads;
+                if (std::regex_match(tensor_name, pattern_qkv_gate_weight)) {
+                    GGML_ASSERT(tensor->ne[axis] == 2*key_dim + 2*value_dim);
+                    return {{key_dim, 2 + 2*head_ratio}};
+                }
+                if (std::regex_match(tensor_name, pattern_ssm_beta_alpha_cat)) {
+                    GGML_ASSERT(tensor->ne[axis] == 2*n_v_heads);
+                    return {{n_k_heads, 2*head_ratio}};
+                }
                 if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_ssm_conv1d)) {
                     GGML_ASSERT(tensor->ne[axis] == 2*key_dim + value_dim);
                     return {{key_dim, 2 + head_ratio}};
@@ -811,11 +821,13 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
             const int64_t blck_size_perf  = std::lcm(blck_size, 128);
             const int64_t granularity_qkv = std::lcm(blck_size_perf, head_dim);
             if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_attn_gate_weight) ||
-                    std::regex_match(tensor_name, pattern_ssm_conv1d) || std::regex_match(tensor_name, pattern_ssm_out_weight)) {
+                    std::regex_match(tensor_name, pattern_ssm_conv1d) || std::regex_match(tensor_name, pattern_ssm_out_weight) ||
+                    std::regex_match(tensor_name, pattern_qkv_gate_weight)) {
                 return std::vector<int64_t>(segments.size(), granularity_qkv);
             }
             if (std::regex_match(tensor_name, pattern_ssm_dt) || std::regex_match(tensor_name, pattern_ssm_a) ||
-                    std::regex_match(tensor_name, pattern_ssm_alpha) || std::regex_match(tensor_name, pattern_ssm_beta)) {
+                    std::regex_match(tensor_name, pattern_ssm_alpha) || std::regex_match(tensor_name, pattern_ssm_beta) ||
+                    std::regex_match(tensor_name, pattern_ssm_beta_alpha_cat)) {
                 return std::vector<int64_t>(segments.size(), granularity_qkv / head_dim);
             }
             if (std::regex_match(tensor_name, pattern_ssm_beta_alpha)) {
@@ -3566,10 +3578,6 @@ const llama_vocab * llama_model_get_vocab(const llama_model * model) {
     return &model->vocab;
 }
 
-void llama_free_model(llama_model * model) {
-    llama_model_free(model);
-}
-
 void llama_model_free(llama_model * model) {
     delete model;
 }
@@ -3640,26 +3648,6 @@ const char * llama_model_decision_label(const struct llama_model * model, uint32
     return nullptr;
 }
 
-// deprecated
-int32_t llama_n_ctx_train(const llama_model * model) {
-    return llama_model_n_ctx_train(model);
-}
-
-// deprecated
-int32_t llama_n_embd(const llama_model * model) {
-    return llama_model_n_embd(model);
-}
-
-// deprecated
-int32_t llama_n_layer(const llama_model * model) {
-    return llama_model_n_layer(model);
-}
-
-// deprecated
-int32_t llama_n_head(const llama_model * model) {
-    return llama_model_n_head(model);
-}
-
 llama_rope_type llama_model_rope_type(const llama_model * model) {
     switch (model->arch) {
         // these models do not use RoPE

src/llama-sampler.cpp

@@ -2697,8 +2697,6 @@ static struct llama_sampler * llama_sampler_init_grammar_impl(
                       const char * grammar_str,
                       const char * grammar_root,
                               bool lazy,
-                     const char ** trigger_words,
-                            size_t num_trigger_words,
                const llama_token * trigger_tokens,
                             size_t num_trigger_tokens,
                      const char ** trigger_patterns,
@@ -2727,7 +2725,7 @@ static void llama_sampler_grammar_reset(struct llama_sampler * smpl) {
 static struct llama_sampler * llama_sampler_grammar_clone(const struct llama_sampler * smpl) {
     const auto * ctx = (const llama_sampler_grammar *) smpl->ctx;
 
-    auto * result = llama_sampler_init_grammar_impl(ctx->vocab, nullptr, nullptr, false, nullptr, 0, nullptr, 0, nullptr, 0);
+    auto * result = llama_sampler_init_grammar_impl(ctx->vocab, nullptr, nullptr, false, nullptr, 0, nullptr, 0);
     GGML_ASSERT(result);
 
     // copy the state
@@ -2775,8 +2773,6 @@ static struct llama_sampler * llama_sampler_init_grammar_impl(
                       const char * grammar_str,
                       const char * grammar_root,
                               bool lazy,
-                     const char ** trigger_words,
-                            size_t num_trigger_words,
                const llama_token * trigger_tokens,
                             size_t num_trigger_tokens,
                      const char ** trigger_patterns,
@@ -2784,26 +2780,7 @@ static struct llama_sampler * llama_sampler_init_grammar_impl(
     auto * ctx = new llama_sampler_grammar;
 
     if (grammar_str != nullptr && grammar_str[0] != '\0') {
-        std::string trigger_pattern;
-        llama_grammar * grammar = nullptr;
-        // TODO: remove trigger_words support.
-        if (trigger_words != nullptr && num_trigger_words > 0) {
-            GGML_ASSERT(trigger_patterns == nullptr && num_trigger_patterns == 0);
-            trigger_pattern = "[\\s\\S]*?(";
-            for (size_t i = 0; i < num_trigger_words; ++i) {
-                static const std::regex special_chars("[.^$|()*+?\\[\\]{}\\\\]");
-                if (i > 0) {
-                    trigger_pattern += "|";
-                }
-                trigger_pattern += std::regex_replace(trigger_words[i], special_chars, "\\$0");
-            }
-            trigger_pattern += ")[\\s\\S]*";
-
-            std::array<const char *, 1> tmp_trigger_patterns = { trigger_pattern.c_str() };
-            grammar = llama_grammar_init_impl(vocab, grammar_str, grammar_root, lazy, tmp_trigger_patterns.data(), tmp_trigger_patterns.size(), trigger_tokens, num_trigger_tokens);
-        } else {
-            grammar = llama_grammar_init_impl(vocab, grammar_str, grammar_root, lazy, trigger_patterns, num_trigger_patterns, trigger_tokens, num_trigger_tokens);
-        }
+        llama_grammar * grammar = llama_grammar_init_impl(vocab, grammar_str, grammar_root, lazy, trigger_patterns, num_trigger_patterns, trigger_tokens, num_trigger_tokens);
         *ctx = {
             /* .vocab        = */ vocab,
             /* .grammar_str  = */ grammar_str,
@@ -2833,18 +2810,7 @@ struct llama_sampler * llama_sampler_init_grammar(
         const struct llama_vocab * vocab,
                       const char * grammar_str,
                       const char * grammar_root) {
-    return llama_sampler_init_grammar_impl(vocab, grammar_str, grammar_root, /* lazy= */ false, nullptr, 0, nullptr, 0, nullptr, 0);
-}
-
-struct llama_sampler * llama_sampler_init_grammar_lazy(
-        const struct llama_vocab * vocab,
-                      const char * grammar_str,
-                      const char * grammar_root,
-                     const char ** trigger_words,
-                            size_t num_trigger_words,
-               const llama_token * trigger_tokens,
-                            size_t num_trigger_tokens) {
-    return llama_sampler_init_grammar_impl(vocab, grammar_str, grammar_root, /* lazy= */ true, trigger_words, num_trigger_words, trigger_tokens, num_trigger_tokens, nullptr, 0);
+    return llama_sampler_init_grammar_impl(vocab, grammar_str, grammar_root, /* lazy= */ false, nullptr, 0, nullptr, 0);
 }
 
 struct llama_sampler * llama_sampler_init_grammar_lazy_patterns(
@@ -2855,7 +2821,7 @@ struct llama_sampler * llama_sampler_init_grammar_lazy_patterns(
                             size_t num_trigger_patterns,
                const llama_token * trigger_tokens,
                             size_t num_trigger_tokens) {
-    return llama_sampler_init_grammar_impl(vocab, grammar_str, grammar_root, /* lazy= */ true, nullptr, 0, trigger_tokens, num_trigger_tokens, trigger_patterns, num_trigger_patterns);
+    return llama_sampler_init_grammar_impl(vocab, grammar_str, grammar_root, /* lazy= */ true, trigger_tokens, num_trigger_tokens, trigger_patterns, num_trigger_patterns);
 }
 
 // penalties