graph: store the conv state of every rollback arch in one set_rows

The rollback slots of the causal conv history were written with one
copy per slot in qwen4exp, glm5-next, kimi-k3, bailingmoe3, lfm and
mamba2, while delta_net already wrote them with a single set_rows.
build_conv_state_store now does this for all of them: one copy without
rollback, otherwise one set_rows of the overlapping windows into the
min(n, n_rs_seq + 1) slots a rollback can read.

The slot order {n_rs_seq, ..., 0} is a constant of the recurrent
memory, kept in its own buffer next to the state, and each graph takes
a view of its tail. Its size no longer depends on the ubatch, so the
reserved graph covers every ubatch shape, and the graph input that
held it is gone.

A model that is not causal never rolls back, so the context clamps
n_rs_seq to 0 for it, like for an arch without rollback support.

.github/workflows/pr-video-demo-label.yml

@@ -0,0 +1,37 @@
+name: "PR video demo label"
+
+on:
+  pull_request_target:
+    types: [opened, edited, reopened, synchronize]
+
+cache-mode: none
+permissions:
+  pull-requests: write
+  issues: write
+
+jobs:
+  video-demo-label:
+    runs-on: ubuntu-slim
+    steps:
+      - name: Add or remove the video demo label
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          PR_URL: ${{ github.event.pull_request.html_url }}
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          # passed through env, so the body is only matched, never run, and is safe to read from a fork
+          BODY: ${{ github.event.pull_request.body }}
+          HAS_LABEL: ${{ contains(github.event.pull_request.labels.*.name, 'video demo') }}
+        run: |
+          # sections written by demo:pr (see tools/ui/demo/cli.ts): the managed comment
+          # markers, or a preview linking to a demo video uploaded for this PR
+          has_demo=false
+          if { grep -qF '<!-- demo:start -->' <<< "$BODY" && grep -qF '<!-- demo:end -->' <<< "$BODY"; } ||
+            grep -qE "/resolve/videos/pr/$PR_NUMBER/[0-9a-f]{10}(-[0-9]+)?\.mp4" <<< "$BODY"; then
+            has_demo=true
+          fi
+
+          if [ "$has_demo" = true ] && [ "$HAS_LABEL" = false ]; then
+            gh pr edit "$PR_URL" --add-label 'video demo'
+          elif [ "$has_demo" = false ] && [ "$HAS_LABEL" = true ]; then
+            gh pr edit "$PR_URL" --remove-label 'video demo'
+          fi

README.md

@@ -13,7 +13,7 @@
 [![Docker](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/docker.yml?label=Docker)](https://github.com/ggml-org/llama.cpp/actions/workflows/docker.yml)
 [![Winget](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/winget.yml?label=Winget)](https://github.com/ggml-org/llama.cpp/actions/workflows/winget.yml)
 
-[ggml](https://github.com/ggml-org/ggml) / [ops](https://github.com/ggml-org/llama.cpp/blob/master/docs/ops.md) / [maintainer PRs](https://github.com/ggml-org/llama.cpp/issues?q=is%3Apr%20is%3Aopen%20draft%3AFalse%20(author%3Asanmai%20OR%20author%3Argerganov%20OR%20author%3AKitaitiMakoto%20OR%20author%3Adanbev%20OR%20author%3Aaldehir%20OR%20author%3Amax-krasnyansky%20OR%20author%3ACISC%20OR%20author%3Aggerganov%20OR%20author%3Aam17an%20OR%20author%3Ajhen0409%20OR%20author%3Abartowski1182%20OR%20author%3Anikwen%20OR%20author%3Ahipudding%20OR%20author%3Aravi9%20OR%20author%3AServeurpersoCom%20OR%20author%3Athelittlefireman%20OR%20author%3Apwilkin%20OR%20author%3Areeselevine%20OR%20author%3Angxson%20OR%20author%3Ajeffbolznv%20OR%20author%3Amarty1885%20OR%20author%3A0cc4m%20OR%20author%3ATitaniumtown%20OR%20author%3Aangt%20OR%20author%3AIMbackK%20OR%20author%3Aarthw%20OR%20author%3AJohannesGaessler%20OR%20author%3AORippler%20OR%20author%3Aruixiang63%20OR%20author%3Axctan%20OR%20author%3AmasterFoad%20OR%20author%3Aallozaur%20OR%20author%3AAndreas-Krebbel%20OR%20author%3Ayomaytk%20OR%20author%3Aaendk%20OR%20author%3Awine99%20OR%20author%3Agaugarg-nv%20OR%20author%3Ataronaeo%20OR%20author%3Aforforever73%20OR%20author%3Achaxu01%20OR%20author%3Alhez%20OR%20author%3AAlekseiNikiforovIBM%20OR%20author%3Anetrunnereve%20OR%20author%3Afairydreaming%20OR%20author%3Aynankani)%20sort%3Aupdated-desc) / [dev stats](https://github.com/ggml-org/llama.cpp-dev) / [lib llama API](https://github.com/ggml-org/llama.cpp/issues/9289) / [llama-server REST API](https://github.com/ggml-org/llama.cpp/issues/9291)
+[ggml](https://github.com/ggml-org/ggml) / [ops](https://github.com/ggml-org/llama.cpp/blob/master/docs/ops.md) / [maintainer PRs](https://github.com/ggml-org/llama.cpp/pulls?q=is%3Apr%20is%3Aopen%20draft%3AFalse%20(author%3Asanmai%20OR%20author%3Argerganov%20OR%20author%3AKitaitiMakoto%20OR%20author%3Adanbev%20OR%20author%3Aaldehir%20OR%20author%3Amax-krasnyansky%20OR%20author%3ACISC%20OR%20author%3Aggerganov%20OR%20author%3Aam17an%20OR%20author%3Ajhen0409%20OR%20author%3Abartowski1182%20OR%20author%3Anikwen%20OR%20author%3Ahipudding%20OR%20author%3Aravi9%20OR%20author%3AServeurpersoCom%20OR%20author%3Athelittlefireman%20OR%20author%3Apwilkin%20OR%20author%3Areeselevine%20OR%20author%3Angxson%20OR%20author%3Ajeffbolznv%20OR%20author%3Amarty1885%20OR%20author%3A0cc4m%20OR%20author%3ATitaniumtown%20OR%20author%3Aangt%20OR%20author%3AIMbackK%20OR%20author%3Aarthw%20OR%20author%3AJohannesGaessler%20OR%20author%3AORippler%20OR%20author%3Aruixiang63%20OR%20author%3Axctan%20OR%20author%3AmasterFoad%20OR%20author%3Aallozaur%20OR%20author%3AAndreas-Krebbel%20OR%20author%3Ayomaytk%20OR%20author%3Aaendk%20OR%20author%3Awine99%20OR%20author%3Agaugarg-nv%20OR%20author%3Ataronaeo%20OR%20author%3Aforforever73%20OR%20author%3Achaxu01%20OR%20author%3Alhez%20OR%20author%3AAlekseiNikiforovIBM%20OR%20author%3Anetrunnereve%20OR%20author%3Afairydreaming%20OR%20author%3Aynankani)%20sort%3Aupdated-desc) / [dev stats](https://github.com/ggml-org/llama.cpp-dev) / [lib llama API](https://github.com/ggml-org/llama.cpp/issues/9289) / [llama-server REST API](https://github.com/ggml-org/llama.cpp/issues/9291)
 
 </div>
 
@@ -79,7 +79,7 @@ The `llama.cpp` project is build on top of the [ggml](https://github.com/ggml-or
 | --- | --- |
 | [BLAS](docs/build.md#blas-build) | All |
 | [BLIS](docs/backend/BLIS.md) | All |
-| [CANN](docs/build.md#cann) | Ascend NPU |
+| [CANN](docs/build.md#cann) | Huawei Ascend NPU |
 | [CUDA](docs/build.md#cuda) | Nvidia GPU |
 | [HIP](docs/build.md#hip) | AMD GPU |
 | [Hexagon](docs/backend/snapdragon/README.md) | Snapdragon |

common/common.cpp

@@ -2275,6 +2275,16 @@ int32_t common_batch::add_embd(llama_embd embd, const llama_pos * pos, llama_seq
     return size() - 1;
 }
 
+int32_t common_batch::add_embd_row(int32_t row, const llama_pos * pos, llama_seq_id seq_id, bool output) {
+    token t = { LLAMA_TOKEN_NULL, { 0, 0, 0, 0 }, seq_id, output, { nullptr, 0, 0 }, { nullptr, 0, 0 }, {} };
+    t.embd_row = row;
+    for (int32_t j = 0; j < n_pos; ++j) {
+        t.pos[j] = pos[j];
+    }
+    tokens.push_back(t);
+    return size() - 1;
+}
+
 llama_batch_ext * common_batch::get_sub_batch(int32_t off, int32_t n) {
     GGML_ASSERT(batch && "common_batch was not initialized with a context");
     GGML_ASSERT(off >= 0 && n >= 0 && off + n <= size());
@@ -2295,6 +2305,12 @@ llama_batch_ext * common_batch::get_sub_batch(int32_t off, int32_t n) {
             if (t.embd.data && !llama_batch_ext_set_embd_token(res, idx, t.embd)) {
                 GGML_ABORT("%s: failed to set the embedding of token %d at index %d\n", __func__, t.id, i);
             }
+        } else if (t.embd_row >= 0) {
+            idx = llama_batch_ext_add_embd_row(res, t.seq_id, t.embd_row);
+            if (idx < 0) {
+                GGML_ABORT("%s: failed to add embd row %d at index %d (error %d, n = %d)\n", __func__, t.embd_row, i, idx, n);
+            }
+            llama_batch_ext_set_pos(res, idx, t.pos.data());
         } else {
             idx = llama_batch_ext_add_embd(res, t.seq_id, t.embd);
             if (idx < 0) {

common/common.h

@@ -1096,6 +1096,7 @@ struct common_batch {
         llama_embd   state; // non-owning view of the data passed to set_embd_state(), data == NULL if none
         std::vector<llama_seq_id> seq_ids_extra; // see add_seq()
         int32_t      decision_order = 0; // see llama_batch_ext_set_decision_order()
+        int32_t      embd_row = -1; // see add_embd_row(), -1 if none
     };
 
     std::vector<token> tokens; // mirror of the entries, tokens[i] describes batch index i
@@ -1138,6 +1139,10 @@ struct common_batch {
     // pos points to n_pos positions
     int32_t add_embd(llama_embd embd, const llama_pos * pos, llama_seq_id seq_id, bool output);
 
+    // add an entry whose embedding is row `row` of the layer input device buffer of the context's ctx_other
+    // pos points to n_pos positions
+    int32_t add_embd_row(int32_t row, const llama_pos * pos, llama_seq_id seq_id, bool output);
+
     int32_t size() const { return (int32_t) tokens.size(); }
 };
 

common/jinja/runtime.cpp

@@ -307,58 +307,13 @@ value binary_expression::execute_impl(context & ctx) const {
     throw std::runtime_error("Unknown operator \"" + op.value + "\" between " + left_val->type() + " and " + right_val->type());
 }
 
-static value try_builtin_func(context & ctx, const std::string & name, value & input, bool undef_on_missing = false) {
-    JJ_DEBUG("Trying built-in function '%s' for type %s", name.c_str(), input->type().c_str());
-    if (ctx.is_get_stats) {
-        value_t::stats_t::mark_used(input);
-        input->stats.ops.insert(name);
-    }
-    auto builtins = input->get_builtins();
-    auto it = builtins.find(name);
-    if (it != builtins.end()) {
-        JJ_DEBUG("Binding built-in '%s'", name.c_str());
-        return mk_val<value_func>(name, it->second, input);
-    }
-    if (undef_on_missing) {
-        return mk_val<value_undefined>(name);
-    }
-    throw std::runtime_error("Unknown (built-in) filter '" + name + "' for type " + input->type());
-}
-
 static value apply_filter(context & ctx, const statement_ptr & filter, value input) {
     JJ_DEBUG("Applying filter to %s", input->type().c_str());
 
-    auto set_filter_alias = [](auto & filter_id) {
-        if (filter_id == "count") {
-            filter_id = "length";
-        } else if (filter_id == "d") {
-            filter_id = "default";
-        } else if (filter_id == "e") {
-            filter_id = "escape";
-        } else if (filter_id == "trim") {
-            filter_id = "strip";
-        }
-    };
-
     if (is_stmt<identifier>(filter)) {
         auto filter_id = cast_stmt<identifier>(filter)->val;
-
-        set_filter_alias(filter_id);
         JJ_DEBUG("Applying filter '%s' to %s", filter_id.c_str(), input->type().c_str());
-        // TODO: Refactor filters so this coercion can be done automatically
-        if (!input->is_undefined() && !is_val<value_string>(input) && (
-            filter_id == "capitalize" ||
-            filter_id == "lower" ||
-            filter_id == "replace" ||
-            filter_id == "strip" ||
-            filter_id == "title" ||
-            filter_id == "upper" ||
-            filter_id == "wordcount"
-        )) {
-            JJ_DEBUG("Coercing %s to String for '%s' filter", input->type().c_str(), filter_id.c_str());
-            input = mk_val<value_string>(input->as_string());
-        }
-        return try_builtin_func(ctx, filter_id, input)->invoke(func_args(ctx));
+        return apply_filter_by_name(ctx, filter_id, input, func_args(ctx));
 
     } else if (is_stmt<call_expression>(filter)) {
         auto call = cast_stmt<call_expression>(filter);
@@ -367,14 +322,13 @@ static value apply_filter(context & ctx, const statement_ptr & filter, value inp
         }
         auto filter_id = cast_stmt<identifier>(call->callee)->val;
 
-        set_filter_alias(filter_id);
         JJ_DEBUG("Applying filter '%s' with arguments to %s", filter_id.c_str(), input->type().c_str());
         func_args args(ctx);
         for (const auto & arg_expr : call->args) {
             args.push_back(arg_expr->execute(ctx));
         }
 
-        return try_builtin_func(ctx, filter_id, input)->invoke(args);
+        return apply_filter_by_name(ctx, filter_id, input, args);
 
     } else {
         throw std::runtime_error("Invalid filter expression");

common/jinja/value.cpp

@@ -72,6 +72,17 @@ const std::vector<value> & func_args::get_args() const {
     return args;
 }
 
+void func_args::ensure_kwargs(std::initializer_list<const char *> names) const {
+    for (const auto & arg : args) {
+        if (is_val<value_kwarg>(arg)) {
+            auto * kwarg = cast_val<value_kwarg>(arg);
+            if (std::find(names.begin(), names.end(), kwarg->key) == names.end()) {
+                throw raised_exception("Function '" + func_name + "' got an unexpected keyword argument '" + kwarg->key + "'");
+            }
+        }
+    }
+}
+
 /**
  * Function that mimics Python's array slicing.
  */
@@ -134,6 +145,14 @@ static value empty_value_fn(const func_args &) {
         return mk_val<T>();
     }
 }
+
+template<typename T>
+static value empty_value_no_args_fn(const func_args & args) {
+    args.ensure_kwargs();
+    args.ensure_count(1, 1);
+    return empty_value_fn<T>(args);
+}
+
 template<typename T>
 static value test_type_fn(const func_args & args) {
     args.ensure_count(1);
@@ -594,6 +613,52 @@ const func_builtins & global_builtins() {
     return builtins;
 }
 
+value try_builtin_func(context & ctx, const std::string & name, value & input, bool undef_on_missing) {
+    JJ_DEBUG("Trying built-in function '%s' for type %s", name.c_str(), input->type().c_str());
+    if (ctx.is_get_stats) {
+        value_t::stats_t::mark_used(input);
+        input->stats.ops.insert(name);
+    }
+    auto builtins = input->get_builtins();
+    auto it = builtins.find(name);
+    if (it != builtins.end()) {
+        JJ_DEBUG("Binding built-in '%s'", name.c_str());
+        return mk_val<value_func>(name, it->second, input);
+    }
+    if (undef_on_missing) {
+        return mk_val<value_undefined>(name);
+    }
+    throw std::runtime_error("Unknown (built-in) filter '" + name + "' for type " + input->type());
+}
+
+value apply_filter_by_name(context & ctx, std::string filter_id, value input, const func_args & args) {
+    if (filter_id == "count") {
+        filter_id = "length";
+    } else if (filter_id == "d") {
+        filter_id = "default";
+    } else if (filter_id == "e") {
+        filter_id = "escape";
+    } else if (filter_id == "trim") {
+        filter_id = "strip";
+    }
+
+    // TODO: Refactor filters so this coercion can be done automatically
+    if (!input->is_undefined() && !is_val<value_string>(input) && (
+        filter_id == "capitalize" ||
+        filter_id == "lower" ||
+        filter_id == "replace" ||
+        filter_id == "strip" ||
+        filter_id == "title" ||
+        filter_id == "upper" ||
+        filter_id == "wordcount"
+    )) {
+        JJ_DEBUG("Coercing %s to String for '%s' filter", input->type().c_str(), filter_id.c_str());
+        input = mk_val<value_string>(input->as_string());
+    }
+
+    return try_builtin_func(ctx, filter_id, input)->invoke(args);
+}
+
 
 const func_builtins & value_int_t::get_builtins() const {
     static const func_builtins builtins = {
@@ -667,16 +732,22 @@ const func_builtins & value_string_t::get_builtins() const {
     static const func_builtins builtins = {
         {"default", default_value},
         {"upper", [](const func_args & args) -> value {
+            args.ensure_kwargs();
+            args.ensure_count(1, 1);
             args.ensure_vals<value_string>();
             jinja::string str = args.get_pos(0)->as_string().uppercase();
             return mk_val<value_string>(str);
         }},
         {"lower", [](const func_args & args) -> value {
+            args.ensure_kwargs();
+            args.ensure_count(1, 1);
             args.ensure_vals<value_string>();
             jinja::string str = args.get_pos(0)->as_string().lowercase();
             return mk_val<value_string>(str);
         }},
         {"strip", [](const func_args & args) -> value {
+            args.ensure_kwargs({"chars"});
+            args.ensure_count(1, 2);
             value val_input = args.get_pos(0);
             if (!is_val<value_string>(val_input)) {
                 throw raised_exception("strip() first argument must be a string");
@@ -707,11 +778,15 @@ const func_builtins & value_string_t::get_builtins() const {
             }
         }},
         {"title", [](const func_args & args) -> value {
+            args.ensure_kwargs();
+            args.ensure_count(1, 1);
             args.ensure_vals<value_string>();
             jinja::string str = args.get_pos(0)->as_string().titlecase();
             return mk_val<value_string>(str);
         }},
         {"capitalize", [](const func_args & args) -> value {
+            args.ensure_kwargs();
+            args.ensure_count(1, 1);
             args.ensure_vals<value_string>();
             jinja::string str = args.get_pos(0)->as_string().capitalize();
             return mk_val<value_string>(str);
@@ -1013,6 +1088,10 @@ const func_builtins & value_bool_t::get_builtins() const {
     return builtins;
 }
 
+[[noreturn]] static value array_sum_not_implemented(const func_args &) {
+    throw not_implemented_exception("Array sum builtin not implemented");
+}
+
 const func_builtins & value_array_t::get_builtins() const {
     static const func_builtins builtins = {
         {"default", default_value},
@@ -1079,7 +1158,9 @@ const func_builtins & value_array_t::get_builtins() const {
         {"select", select_or_reject<false, false>},
         {"rejectattr", select_or_reject<true, true>},
         {"reject", select_or_reject<true, false>},
+        {"sum", array_sum_not_implemented},
         {"join", [](const func_args & args) -> value {
+            args.ensure_kwargs({"d", "attribute"});
             args.ensure_count(1, 3);
             if (!is_val<value_array>(args.get_pos(0))) {
                 throw raised_exception("join() first argument must be an array");
@@ -1117,13 +1198,28 @@ const func_builtins & value_array_t::get_builtins() const {
         {"tojson", tojson},
         {"map", [](const func_args & args) -> value {
             args.ensure_count(2);
-            if (!is_val<value_array>(args.get_pos(0))) {
+            value val = args.get_pos(0);
+            if (!is_val<value_array>(val)) {
                 throw raised_exception("map: first argument must be an array");
             }
             if (!is_val<value_kwarg>(args.get_args().at(1))) {
-                throw not_implemented_exception("map: filter-mapping not implemented");
+                value filter = args.get_pos(1);
+                if (!is_val<value_string>(filter)) {
+                    throw raised_exception("map: filter name must be a string");
+                }
+                auto out = mk_val<value_array>();
+                func_args filter_args(args.ctx);
+                const auto & all_args = args.get_args();
+                for (size_t i = 2; i < all_args.size(); ++i) {
+                    filter_args.push_back(all_args[i]);
+                }
+                const std::string filter_id = filter->as_string().str();
+                auto arr = val->as_array();
+                for (const auto & item : arr) {
+                    out->push_back(apply_filter_by_name(args.ctx, filter_id, item, filter_args));
+                }
+                return is_val<value_tuple>(val) ? mk_val<value_tuple>(out->as_array()) : out;
             }
-            value val       = args.get_pos(0);
             value attribute = args.get_kwarg_or_pos("attribute", 1);
             value default_val = args.get_kwarg("default", mk_val<value_undefined>());
             auto out = mk_val<value_array>();
@@ -1132,7 +1228,7 @@ const func_builtins & value_array_t::get_builtins() const {
                 value attr_val = get_attribute(item, attribute, default_val);
                 out->push_back(attr_val);
             }
-            return is_val<value_tuple>(val) ? mk_val<value_tuple>(std::move(out->as_array())) : out;
+            return is_val<value_tuple>(val) ? mk_val<value_tuple>(out->as_array()) : out;
         }},
         {"append", [](const func_args & args) -> value {
             args.ensure_count(2);
@@ -1405,14 +1501,18 @@ const func_builtins & value_none_t::get_builtins() const {
 const func_builtins & value_undefined_t::get_builtins() const {
     static const func_builtins builtins = {
         {"default", default_value},
-        {"capitalize", empty_value_fn<value_string>},
+        {"capitalize", empty_value_no_args_fn<value_string>},
         {"first", empty_value_fn<value_undefined>},
         {"items", empty_value_fn<value_array>},
-        {"join", empty_value_fn<value_string>},
+        {"join", [](const func_args & args) -> value {
+            args.ensure_kwargs({"d", "attribute"});
+            args.ensure_count(1, 3);
+            return empty_value_fn<value_string>(args);
+        }},
         {"last", empty_value_fn<value_undefined>},
         {"length", empty_value_fn<value_int>},
         {"list", empty_value_fn<value_array>},
-        {"lower", empty_value_fn<value_string>},
+        {"lower", empty_value_no_args_fn<value_string>},
         {"map", empty_value_fn<value_array>},
         {"max", empty_value_fn<value_undefined>},
         {"min", empty_value_fn<value_undefined>},
@@ -1425,12 +1525,16 @@ const func_builtins & value_undefined_t::get_builtins() const {
         {"selectattr", empty_value_fn<value_array>},
         {"sort", empty_value_fn<value_array>},
         {"string", empty_value_fn<value_string>},
-        {"strip", empty_value_fn<value_string>},
+        {"strip", [](const func_args & args) -> value {
+            args.ensure_kwargs({"chars"});
+            args.ensure_count(1, 2);
+            return empty_value_fn<value_string>(args);
+        }},
         {"sum", empty_value_fn<value_int>},
-        {"title", empty_value_fn<value_string>},
+        {"title", empty_value_no_args_fn<value_string>},
         {"truncate", empty_value_fn<value_string>},
         {"unique", empty_value_fn<value_array>},
-        {"upper", empty_value_fn<value_string>},
+        {"upper", empty_value_no_args_fn<value_string>},
         {"wordcount", empty_value_fn<value_int>},
     };
     return builtins;

common/jinja/value.h

@@ -7,6 +7,7 @@
 #include <cmath>
 #include <cstdint>
 #include <functional>
+#include <initializer_list>
 #include <map>
 #include <memory>
 #include <set>
@@ -176,6 +177,8 @@ protected:
 //
 
 const func_builtins & global_builtins();
+value try_builtin_func(context & ctx, const std::string & name, value & input, bool undef_on_missing = false);
+value apply_filter_by_name(context & ctx, std::string filter_id, value input, const func_args & args);
 
 std::string value_to_json(const value & val, int indent = -1, const std::string_view item_sep = ", ", const std::string_view key_sep = ": ");
 
@@ -658,6 +661,7 @@ public:
     size_t count() const { return args.size(); }
     void push_back(const value & val);
     void push_front(const value & val);
+    void ensure_kwargs(std::initializer_list<const char *> names = {}) const;
     void ensure_count(size_t min, size_t max = 999) const {
         size_t n = args.size();
         if (n < min || n > max) {

common/parsers/gpt-oss.cpp

@@ -71,6 +71,7 @@ common_chat_params common_chat_params_init_gpt_oss(const common_chat_template &
         auto content         = p.rule("message-content", p.until("<|end|>"));
         auto channel         = p.literal("<|channel|>") + (p.literal("commentary") | p.literal("analysis"));
         auto constrain_type  = p.chars("[A-Za-z0-9_-]", 1, -1);
+        auto constraint      = p.optional(p.space() + p.optional(p.literal("<|constrain|>")) + constrain_type);
 
         // Occasionally, gpt-oss-20b will prefix channels with this commentary
         auto stray_commentary = p.optional(p.literal("<|channel|>commentary") + p.optional(p.literal(" to=assistant")));
@@ -84,7 +85,7 @@ common_chat_params common_chat_params_init_gpt_oss(const common_chat_template &
 
         auto analysis = p.ref("analysis");
         auto preamble = p.rule("preamble", p.literal("<|channel|>commentary<|message|>") + p.content(content) + end);
-        auto final_msg = p.rule("final", stray_commentary + p.literal("<|channel|>final<|message|>") + p.content(content));
+        auto final_msg = p.rule("final", stray_commentary + p.literal("<|channel|>final") + constraint + p.literal("<|message|>") + p.content(content));
 
         // Consume any unsolicited tool calls, e.g. builtin functions
         auto unsolicited = p.rule("unsolicited", p.atomic(p.optional(channel) + p.literal(" to=") + content + end));
@@ -92,7 +93,6 @@ common_chat_params common_chat_params_init_gpt_oss(const common_chat_template &
         auto any = p.rule("any", preamble | analysis);
 
         if (has_response_format) {
-            auto constraint = p.optional(p.space() + p.optional(p.literal("<|constrain|>")) + constrain_type);
             auto response_format = p.rule("response-format",
                 p.literal("<|channel|>final") + constraint + p.literal("<|message|>") +
                 p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema)));
@@ -109,7 +109,6 @@ common_chat_params common_chat_params_init_gpt_oss(const common_chat_template &
                 const auto   params   = common_chat_tool_parameters(function);
 
                 auto func_name  = p.literal(" to=functions.") + p.tool_name(p.literal(name));
-                auto constraint = p.optional(p.space() + p.optional(p.literal("<|constrain|>")) + constrain_type);
                 auto args       = p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", params));
 
                 // recipient in role header

common/speculative.cpp

@@ -1027,6 +1027,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
     const int32_t * target_layer_ids   = nullptr; // model_dft's extract layer indices
     uint32_t        target_layer_ids_n = 0;
 
+    // inject target features from the target's device buffer instead of host copies (env LLAMA_DFLASH_DEV=0 disables)
+    bool use_dev_features = false;
+
     common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq,
             common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)
         : common_speculative_impl(type, n_seq, params.draft.n_max)
@@ -1141,6 +1144,23 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
             llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);
         }
 
+        // the device buffer keeps layers in ascending id order, the fc input expects target_layer_ids order
+        {
+            const char * env = getenv("LLAMA_DFLASH_DEV");
+            const bool ascending = std::adjacent_find(target_layer_ids, target_layer_ids + target_layer_ids_n,
+                    [](int32_t a, int32_t b) { return a >= b; }) == target_layer_ids + target_layer_ids_n;
+
+            use_dev_features = !is_lilicorr && ascending && (env == nullptr || atoi(env) != 0);
+            if (use_dev_features) {
+                llama_set_embeddings_layer_inp_dev(ctx_tgt, true);
+                use_dev_features = llama_supports_embd_row(ctx_dft);
+                if (!use_dev_features) {
+                    llama_set_embeddings_layer_inp_dev(ctx_tgt, false);
+                }
+            }
+            LOG_INF("%s: DFlash target features via %s\n", __func__, use_dev_features ? "device buffer" : "host copies");
+        }
+
         // DFlash2 reads its selector lattice from h_nextn and never consumes raw logits.
         llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2 && !is_lilicorr);
         llama_set_causal_attn(ctx_dft, causal_attn); // DFlash needs non-causal attention unless the model says otherwise
@@ -1233,6 +1253,25 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
             for (int32_t offset = 0; offset < n_rows; offset += n_ubatch) {
                 const int32_t n_chunk = std::min(n_ubatch, n_rows - offset);
 
+                if (use_dev_features) {
+                    // the features are the target batch rows in the target's device buffer
+                    batch_inject.clear();
+                    for (int32_t i = 0; i < n_chunk; ++i) {
+                        const int32_t   row        = i_batch_beg[seq_id] + offset + i;
+                        const llama_pos p          = batch_in.tokens[row].pos[0];
+                        const llama_pos pos_arr[4] = { p, p, p, 0 };
+                        batch_inject.add_embd_row(row, pos_arr, seq_id, false);
+                    }
+
+                    const int32_t rc = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch_inject.get());
+                    if (rc != 0) {
+                        LOG_ERR("%s: llama_process(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
+                                __func__, rc, (int) n_chunk, (int) offset);
+                        return false;
+                    }
+                    continue;
+                }
+
                 // gather target features per extract layer; the fused decode encodes and
                 // injects them into the K/V cache at the target positions
                 features_buf.resize((size_t) n_chunk * n_embd_enc);

conversion/__init__.py

@@ -46,6 +46,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "ChatGLMForConditionalGeneration": "chatglm",
     "ChatGLMModel": "chatglm",
     "ClefModel": "clef",
+    "CloverMoeModelLM": "clover",
     "CodeShellForCausalLM": "codeshell",
     "CogVLMForCausalLM": "cogvlm",
     "Cohere2MoeForCausalLM": "command_r",
@@ -222,6 +223,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "ModernBertForSequenceClassification": "bert",
     "ModernBertModel": "bert",
     "NanbeigeForCausalLM": "nanbeige",
+    "Nemotron3_5TranscribeForConditionalGeneration": "nemotron",
     "NemotronForCausalLM": "nemotron",
     "NemotronHForCausalLM": "nemotron",
     "NemotronHPuzzleForCausalLM": "nemotron",
@@ -379,6 +381,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
     "MiniCPMOcrForConditionalGeneration": "minicpm",
     "Mistral3ForConditionalGeneration": "llava",
     "Mistral3Model": "llava",
+    "Nemotron3_5TranscribeForConditionalGeneration": "nemotron",
     "NemotronH_Omni_Reasoning_V3": "nemotron",
     "NemotronH_Nano_VL_V2": "nemotron",
     "MuseGlimmerForConditionalGeneration": "muse_glimmer",

conversion/base.py

@@ -1277,6 +1277,7 @@ class ModelBase:
         expert_scales: dict[tuple[int, str], list[tuple[int, float]]] = {}
         expert_input_scales: dict[tuple[int, str], list[tuple[int, float]]] = {}
         expert_shapes: dict[tuple[int, str], list[int]] = {}
+        expert_names: dict[tuple[int, str], str] = {}
         n_experts = self.find_hparam(["num_local_experts", "num_experts", "n_routed_experts"], optional=True) or 0
         consumed: list[str] = []
 
@@ -1325,6 +1326,7 @@ class ModelBase:
                     expert_scales[key] = []
                     expert_input_scales[key] = []
                     expert_shapes[key] = shape
+                    expert_names[key] = name
                 expert_blocks[key].append((expert_id, raw.copy()))
                 # Collect per-expert scale2 (scalar per expert)
                 expert_scales[key].append((expert_id, float(scale2.float().sum())))
@@ -1333,13 +1335,13 @@ class ModelBase:
 
                 # Flush when all experts for this (layer, proj) are collected
                 if n_experts > 0 and len(expert_blocks[key]) >= n_experts:
-                    self._flush_nvfp4_experts(key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, bid, proj_type)
+                    self._flush_nvfp4_experts(key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, expert_names, bid, proj_type)
             else:
                 self._repack_nvfp4(name, weight, scale, scale2, input_scale)
 
         # Flush any remaining experts (fallback if n_experts was unknown)
         for bid, proj_type in list(expert_blocks.keys()):
-            self._flush_nvfp4_experts((bid, proj_type), expert_blocks, expert_scales, expert_input_scales, expert_shapes, bid, proj_type)
+            self._flush_nvfp4_experts((bid, proj_type), expert_blocks, expert_scales, expert_input_scales, expert_shapes, expert_names, bid, proj_type)
 
         # Remove consumed tensors so get_tensors/modify_tensors won't see them
         for name in consumed:
@@ -1350,7 +1352,7 @@ class ModelBase:
             if name.endswith((".k_scale", ".v_scale")):
                 del self.model_tensors[name]
 
-    def _flush_nvfp4_experts(self, key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, bid, proj_type):
+    def _flush_nvfp4_experts(self, key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, expert_names, bid, proj_type):
         experts = expert_blocks.pop(key)
         scales = expert_scales.pop(key)
         input_scales = expert_input_scales.pop(key)
@@ -1362,7 +1364,7 @@ class ModelBase:
         new_name = self.map_tensor_name(merged_name)
         logger.info(f"Repacked {new_name} with shape [{len(experts)}, {shape[0]}, {shape[1]}] and quantization NVFP4")
         self.gguf_writer.add_tensor(new_name, merged, raw_dtype=gguf.GGMLQuantizationType.NVFP4)
-        self._tag_prec_a4(merged_name, new_name)
+        self._tag_prec_a4(expert_names.pop(key), new_name)
 
         scales.sort(key=lambda x: x[0])
         self._write_scales_tensor(new_name.replace(".weight", ".scale"), [s[1] for s in scales])
@@ -3384,7 +3386,7 @@ def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> st
     # Kimi-K3's text_config reports "KimiLinearForCausalLM", which is the older
     # Kimi-Linear-48B architecture and cannot load K3 (no attention residuals,
     # latent MoE, situ, ...). Route on the top-level architecture instead.
-    if model_type == ModelType.TEXT and arch in ("NEOChatModel", "StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration"):
+    if model_type == ModelType.TEXT and arch in ("NEOChatModel", "StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration", "Nemotron3_5TranscribeForConditionalGeneration"):
         return arch
 
     # if "architectures" is found in the sub-config, use that instead

conversion/clover.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from typing import Callable, Iterable, TYPE_CHECKING
+
+from .base import ModelBase, gguf
+from .qwen import _LinearAttentionVReorderBase
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+
+@ModelBase.register("CloverMoeModelLM")
+@ModelBase.example("Diffuspace/Clover-35B-A3B-Preview")
+class CloverMoeModel(_LinearAttentionVReorderBase):
+    model_arch = gguf.MODEL_ARCH.CLOVER_MOE
+    supports_mtp_export = False
+
+    def __init__(self, *args, **kwargs):
+        self.no_mtp = True
+        super().__init__(*args, **kwargs)
+
+    def set_vocab(self):
+        self._set_vocab_gpt2()
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_expert_shared_count(self.hparams.get("num_shared_experts", 1))
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, _ = item
+        if name.endswith((".conv_right_weight", ".bd_softmax_logit_scale", ".bd_softmax_history_bias", ".bd_softmax_history_c")):
+            return None
+
+        return super().filter_tensors(item)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name.endswith(".A_log") or (name.endswith("norm.weight") and not name.endswith("linear_attn.norm.weight")):
+            data_torch = data_torch.float()
+
+        yield from super().modify_tensors(data_torch, name, bid)

conversion/nemotron.py

@@ -428,6 +428,9 @@ class NemotronHModel(GraniteHybridModel):
     is_moe: bool = False
     supports_mtp_export = True
     _experts: list[dict[str, Tensor]] | None = None
+    # tokens past the embedding rows (left out of the vocab), and special tokens that are printed as text
+    _vocab_skip: tuple[str, ...] = ()
+    _user_defined_tokens: re.Pattern[str] | None = None
 
     _SSM_LAYER_TYPES = {"mamba", "linear_attention"}
     _ATTN_LAYER_TYPES = {"attention", "full_attention"}
@@ -616,7 +619,11 @@ class NemotronHModel(GraniteHybridModel):
         vocab_size = -(vocab_size // -pad_vocab) * pad_vocab
         self.hparams["vocab_size"] = vocab_size
 
-        assert max(tokenizer.vocab.values()) < vocab_size  # ty: ignore[unresolved-attribute]
+        # draft-model converters (DFlash, LiLiCorr) call this with their own instance, which has neither hook
+        vocab_skip = getattr(self, "_vocab_skip", ())
+        user_defined_tokens = getattr(self, "_user_defined_tokens", None)
+
+        assert max(id_ for tok, id_ in tokenizer.vocab.items() if tok not in vocab_skip) < vocab_size  # ty: ignore[unresolved-attribute]
 
         tokpre = self.get_vocab_base_pre(tokenizer)
 
@@ -638,7 +645,8 @@ class NemotronHModel(GraniteHybridModel):
                         if previous_token != token:
                             logger.info(f"{repr(previous_token)} is encoded and decoded back to {repr(token)} using AutoTokenizer")
 
-                    if added_tokens_decoder[i].special or self.does_token_look_special(token):
+                    printed = user_defined_tokens is not None and user_defined_tokens.fullmatch(token)
+                    if (added_tokens_decoder[i].special or self.does_token_look_special(token)) and not printed:
                         toktypes.append(gguf.TokenType.CONTROL)
                     else:
                         token = token.replace(b"\xe2\x96\x81".decode("utf-8"), " ")  # pre-normalize user-defined spaces
@@ -834,3 +842,140 @@ class NemotronHPuzzleModel(NemotronHModel):
         if item[0].startswith("mtp."):
             return None
         return super().filter_tensors(item)
+
+
+# NeMo speechlm2 SALM checkpoint: a Nemotron-H decoder under llm.*, the speech encoder under perception.*
+@ModelBase.register("Nemotron3_5TranscribeForConditionalGeneration")
+@ModelBase.example("nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B")
+class NemotronTranscribeModel(NemotronHModel):
+    model_arch = gguf.MODEL_ARCH.NEMOTRON_H
+    # the audio placeholder (id vocab_size) has no embedding row; speaker, event and style tags print with the transcript
+    _vocab_skip = ("<|audio|>",)
+    _user_defined_tokens = re.compile(r"<spk:\d+>|<ev:[a-z_]+>|<sty:(?:stt|end):[a-z_]+>")
+    _renames = {
+        "llm.model.embed_tokens.weight": "backbone.embeddings.weight",
+        "llm.model.norm.weight":         "backbone.norm_f.weight",
+        "llm.lm_head.weight":            "lm_head.weight",
+    }
+    # grouped experts: [n_expert, n_in, n_out] -> [n_expert, n_out, n_in]
+    _grouped_experts = {"gate_and_up_projs": "up_proj", "down_projs": "down_proj"}
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+        if name.startswith(("perception.", "ctc_timestamp.")) or "_extra_state" in name:
+            return None
+        if not name.startswith("llm."):
+            # the FP8 and NVFP4 releases store the decoder as backbone.*
+            raise ValueError(f"unexpected tensor {name}: only the BF16 checkpoint layout (llm.*) is supported")
+        if name in cls._renames:
+            name = cls._renames[name]
+        elif name.startswith("llm.model.layers."):
+            name = "backbone.layers." + name[len("llm.model.layers."):]
+        elif name.startswith("llm.mtp.layers."):
+            name = "mtp.layers." + name[len("llm.mtp.layers."):]
+        else:
+            raise ValueError(f"unexpected decoder tensor {name}")
+        return super().filter_tensors((name.replace("._fp32_params.", "."), gen))
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        leaf = name.rsplit(".", 1)[-1]
+        if ".mixer.experts." in name and leaf in self._grouped_experts:
+            if name.startswith("mtp."):
+                assert self._mtp_bid is not None, f"{name}: MTP head is not exported"
+                bid = self._mtp_bid
+            assert bid is not None
+            data_torch = data_torch.permute(0, 2, 1).contiguous()
+            name = f"model.layers.{bid}.mlp.experts.{self._grouped_experts[leaf]}.weight"
+            yield from ModelBase.modify_tensors(self, data_torch, name, bid)
+            return
+        yield from super().modify_tensors(data_torch, name, bid)
+
+
+@ModelBase.register("Nemotron3_5TranscribeForConditionalGeneration")
+class NemotronTranscribeAudioModel(MmprojModel):
+    # NeMo ParallelExpertEncoder: ASR and speaker transformers, speaker activity fused into the ASR states
+    has_vision_encoder = False
+    has_audio_encoder = True
+
+    def get_audio_config(self) -> dict[str, Any] | None:
+        asr = self.global_config["pe_encoder_config"]["asr_encoder_cfg"]
+        return {
+            "hidden_size":         asr["d_model"],
+            "intermediate_size":   int(float(asr.get("ff_expansion", 4)) * asr["d_model"]),
+            "num_hidden_layers":   asr["n_layers"],
+            "num_attention_heads": asr["n_heads"],
+        }
+
+    @staticmethod
+    def _rope_dims(cfg: dict[str, Any]) -> int:
+        assert cfg.get("self_attention_model") == "rope" and cfg.get("subsampling", "feature_stacking") == "feature_stacking"
+        assert cfg.get("attn_mode", "full") == "full" and not cfg.get("xscaling", False) and cfg.get("pre_block_norm", True)
+        assert float(cfg.get("rope_base", 10000.0)) == 10000.0
+        n_rot = int(cfg["d_model"] // cfg["n_heads"] * float(cfg.get("rotary_fraction", 1.0)))
+        assert n_rot >= 2 and n_rot % 2 == 0
+        return n_rot
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        pe = self.global_config["pe_encoder_config"]
+        asr = pe["asr_encoder_cfg"]
+        diar = pe["diarization_model_cfg"]
+        spk = diar["encoder"]
+        assert pe.get("asr_encoder_type") == "transformer" and pe.get("speaker_feature_mode", "thresholded") == "thresholded"
+        assert pe.get("asr_normalize_type") == "per_feature" and pe.get("diar_normalize_type") in (None, "NA")
+        assert asr["subsampling_factor"] == spk["subsampling_factor"] == 8 and asr["feat_in"] == spk["feat_in"]
+        assert self._rope_dims(spk) == spk["d_model"] // spk["n_heads"] and not spk.get("qkv_bias", False)
+        assert diar.get("high_resolution", False)
+        # constants in clip: the ParallelExpertEncoder online settings and the Sortformer streaming config
+        assert (pe.get("online_inference_length", 500), pe.get("chunk_left_context", 50), pe.get("chunk_right_context", 50),
+                pe.get("diar_spkcache_len", 188), pe.get("diar_fifo_len", 40), pe.get("diar_spkcache_update_period", 300),
+                pe.get("speaker_activity_threshold", 0.5)) == (500, 50, 50, 188, 40, 300, 0.5), \
+            "streaming settings differ from the constants in tools/mtmd/models/nemotron-transcribe.cpp"
+        sm = diar["sortformer_modules"]
+        assert (sm["spkcache_sil_frames_per_spk"], sm["pred_score_threshold"], sm["scores_boost_latest"], sm["strong_boost_rate"],
+                sm["weak_boost_rate"], sm["min_pos_scores_rate"], sm["use_learnable_sil_emb"]) == (1, 0.25, 0.05, 0.75, 1.5, 0.5, True), \
+            "speaker cache settings differ from the constants in tools/mtmd/models/nemotron-transcribe.cpp"
+
+        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.NEMOTRON_TRANSCRIBE)
+        self.gguf_writer.add_audio_num_mel_bins(asr["feat_in"])
+        self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
+        self.gguf_writer.add_audio_subsampling_factor(asr["subsampling_factor"])
+        self.gguf_writer.add_audio_rope_dims(self._rope_dims(asr))
+        self.gguf_writer.add_audio_spk_block_count(spk["n_layers"])
+        self.gguf_writer.add_audio_spk_head_count(spk["n_heads"])
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        # ParallelExpertEncoder._build_sinusoid_position_encoding
+        pe = self.global_config["pe_encoder_config"]
+        n_spk = int(pe["diarization_model_cfg"]["sortformer_modules"]["num_spks"])
+        d = int(pe["asr_encoder_cfg"]["d_model"])
+        scale = float(pe.get("spk_kernel_scale", 1.0))
+        pos = torch.arange(n_spk, dtype=torch.float32)[:, None]
+        div = torch.exp(torch.arange(0, d, 2, dtype=torch.float32) * -(np.log(10000.0) / d))
+        kernel = torch.zeros(n_spk, d, dtype=torch.float32)
+        kernel[:, 0::2] = torch.sin(pos * div)
+        kernel[:, 1::2] = torch.cos(pos * div)
+        # stored as [d, n_spk]: the graph multiplies it with the speaker activity directly
+        yield self.format_tensor_name(gguf.MODEL_TENSOR.A_FUSION_SPK_KERNEL, suffix=""), (kernel * scale).T.contiguous()
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.A_SPK_HEAD_UPSAMPLE, None):
+            return gguf.GGMLQuantizationType.F16  # ggml_conv_1d kernel
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+        if not name.startswith("perception.") or "_extra_state" in name:
+            return None
+        # training-only head; the speaker model's unused front end
+        if name.startswith(("perception.encoder.diarization_model.sortformer_modules.hidden_to_spks.",
+                            "perception.encoder.diarization_model.preprocessor.")):
+            return None
+        return name, gen
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name == "perception.preprocessor.featurizer.fb":
+            data_torch = data_torch.squeeze(0)
+        yield from super().modify_tensors(data_torch, name, bid)

conversion/pockettts.py

@@ -5,6 +5,7 @@ from pathlib import Path
 from typing import Any, Iterable, TYPE_CHECKING
 
 import torch
+import torch.nn.functional as F
 
 if TYPE_CHECKING:
     from torch import Tensor
@@ -45,6 +46,102 @@ def _tensor_shapes(dir_model: Path) -> dict[str, tuple[int, ...]]:
         return {name: tuple(part[name].shape) for name in part.keys()}
 
 
+def _has_voice_encoder(dir_model: Path) -> bool:
+    # without voice cloning, the mimi encoder weights are all zero
+    part_names = ModelBase.get_model_part_names(dir_model, "model", ".safetensors")
+    with gguf.utility.SafetensorsLocal(dir_model / part_names[0]) as part:
+        nonzero = {bool(part[name].mmap_bytes().any()) for name in part.keys() if name.startswith("mimi.encoder")}
+    if len(nonzero) != 1:
+        raise ValueError(f"{dir_model}: the mimi encoder is partly zeroed out")
+    return nonzero.pop()
+
+
+def _load_safetensors(path: Path, prefixes: tuple[str, ...]) -> dict[str, Tensor]:
+    dtypes = {"F32": torch.float32, "BF16": torch.bfloat16, "I64": torch.int64}
+    with gguf.utility.SafetensorsLocal(path) as part:
+        return {name: torch.from_numpy(part[name].mmap_bytes()).view(dtypes[part[name].dtype]).reshape(part[name].shape)
+                for name in part.keys() if name.startswith(prefixes)}
+
+
+def _rope_pairs(x: Tensor, n_head: int, rope_theta: float) -> Tensor:
+    # ggml NORM mode rope, x: [n_pos, n_embd]
+    n_pos, n_embd = x.shape
+    head_dim = n_embd // n_head
+    angle = torch.arange(n_pos, dtype=x.dtype)[:, None, None] * \
+        torch.pow(rope_theta, -torch.arange(0, head_dim, 2, dtype=x.dtype) / head_dim)
+    x = x.view(n_pos, n_head, head_dim // 2, 2)
+    return torch.stack([x[..., 0] * angle.cos() - x[..., 1] * angle.sin(),
+                        x[..., 0] * angle.sin() + x[..., 1] * angle.cos()], -1).reshape(n_pos, n_embd)
+
+
+def _voice_prompt_embd(weights: dict[str, Tensor], kv: list[Tensor], n_head: int, rope_theta: float) -> tuple[Tensor, float]:
+    # the voice state only has K/V of each layer (kv[il] is [k/v, pos, n_embd]): layer 0 V gives each frame up to its scale, layer 1 V gives the scale
+    # returns the frames without bos_before_voice, and the max relative K/V error when the frames are prompted again
+    w = {name: t.double() for name, t in weights.items()}
+    kv = [t.double() for t in kv]
+    k0, v0, v1 = kv[0][0], kv[0][1], kv[1][1]
+    n_pos, n_embd = v0.shape
+    head_dim = n_embd // n_head
+
+    def p(il: int, name: str) -> Tensor:
+        return w[f"flow_lm.transformer.layers.{il}.{name}"]
+
+    def norm(x: Tensor, il: int, name: str) -> Tensor:
+        return F.layer_norm(x, (n_embd,), p(il, f"{name}.weight"), p(il, f"{name}.bias"), 1e-5)
+
+    def heads(t: Tensor) -> Tensor:
+        return t.reshape(n_pos, n_head, head_dim).transpose(0, 1)
+
+    def rows(t: Tensor) -> Tensor:
+        return t.transpose(0, 1).reshape(n_pos, n_embd)
+
+    def attn(il: int, h: Tensor, k: Tensor, v: Tensor) -> Tensor:
+        q = _rope_pairs(h @ p(il, "self_attn.in_proj.weight").chunk(3)[0].T, n_head, rope_theta)
+        return rows(F.scaled_dot_product_attention(heads(q), heads(k), heads(v), is_causal=True)) @ p(il, "self_attn.out_proj.weight").T
+
+    def ffn(il: int, x: Tensor) -> Tensor:
+        return F.gelu(norm(x, il, "norm2") @ p(il, "linear1.weight").T) @ p(il, "linear2.weight").T
+
+    wv = p(0, "self_attn.in_proj.weight").chunk(3)[2]
+    h = torch.linalg.solve(wv, v0.T).T                   # norm1(x)
+    z = (h - p(0, "norm1.bias")) / p(0, "norm1.weight")  # (x - mean(x)) / std(x)
+
+    # fit z on [speaker_proj, 1] to get the latent direction back
+    proj = w["flow_lm.speaker_proj_weight"]
+    fit = torch.linalg.lstsq(torch.cat([proj, torch.ones(n_embd, 1, dtype=proj.dtype)], 1), z[1:].T).solution
+    frames = (proj @ fit[:-1]).T
+
+    a0 = attn(0, h, k0, v0)[1:]
+    wv1 = p(1, "self_attn.in_proj.weight").chunk(3)[2]
+
+    def err(scale: Tensor) -> Tensor:  # per frame relative error of layer 1 V
+        x = scale[:, None] * frames + a0
+        x = x + ffn(0, x)
+        return (norm(x, 1, "norm1") @ wv1.T - v1[1:]).norm(dim=1) / v1[1:].norm(dim=1)
+
+    grid = torch.logspace(-3, 3, 49, dtype=torch.float64)
+    best = torch.stack([err(s.expand(n_pos - 1)) for s in grid]).argmin(0)
+    lo, hi = grid[(best - 1).clamp(min=0)], grid[(best + 1).clamp(max=len(grid) - 1)]
+    for _ in range(40):  # golden section search, all frames at once
+        m1, m2 = hi - 0.618 * (hi - lo), lo + 0.618 * (hi - lo)
+        left = err(m1) < err(m2)
+        lo, hi = torch.where(left, lo, m1), torch.where(left, m2, hi)
+    frames = (lo + hi)[:, None] / 2 * frames
+
+    # prompt the frames again and compare K/V with the shipped state
+    x = torch.cat([w["flow_lm.bos_before_voice"].reshape(1, n_embd), frames])
+    drift = 0.0
+    for il, (k_ref, v_ref) in enumerate(kv):
+        h = norm(x, il, "norm1")
+        _, wk, wv = p(il, "self_attn.in_proj.weight").chunk(3)
+        k, v = _rope_pairs(h @ wk.T, n_head, rope_theta), h @ wv.T
+        for a, b in ((k, k_ref), (v, v_ref)):
+            drift = max(drift, ((a - b).norm(dim=1) / b.norm(dim=1)).median().item())
+        x = x + attn(il, h, k, v)
+        x = x + ffn(il, x)
+    return frames, drift
+
+
 @ModelBase.register_hparams_loader(lambda dir_model: "flow_lm.bos_emb" in _tensor_shapes(dir_model))
 def _load_hparams(dir_model: Path) -> dict[str, Any]:
     logger.info("gguf: detected pocket-tts checkpoint, deriving hparams from tensor shapes")
@@ -194,22 +291,27 @@ class PocketTTSMmprojModel(MmprojModel):
         (gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_V),
     )
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.has_audio_encoder = _has_voice_encoder(self.dir_model)
+
     def set_gguf_parameters(self):
         self.gguf_writer.add_file_type(self.ftype)
         assert self.hparams_audio is not None
 
         # voice-prompt encoder: mimi encoder + speaker_proj
-        self.gguf_writer.add_clip_has_audio_encoder(True)
-        # note: the 24kHz sample rate is hardcoded on the clip.cpp side, like the other audio models
-        self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_SPKENC)
-        self.gguf_writer.add_audio_projection_dim(self.n_embd_text)
-        self.gguf_writer.add_audio_block_count(self.hparams_audio["num_hidden_layers"])
-        self.gguf_writer.add_audio_embedding_length(self.hparams_audio["hidden_size"])
-        self.gguf_writer.add_audio_feed_forward_length(self.hparams_audio["intermediate_size"])
-        self.gguf_writer.add_audio_head_count(self.hparams_audio["num_attention_heads"])
-        self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
-        # mimi convolves the waveform directly, it is passed around as a 1-row "mel"
-        self.gguf_writer.add_audio_num_mel_bins(1)
+        if self.has_audio_encoder:
+            self.gguf_writer.add_clip_has_audio_encoder(True)
+            # note: the 24kHz sample rate is hardcoded on the clip.cpp side, like the other audio models
+            self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_SPKENC)
+            self.gguf_writer.add_audio_projection_dim(self.n_embd_text)
+            self.gguf_writer.add_audio_block_count(self.hparams_audio["num_hidden_layers"])
+            self.gguf_writer.add_audio_embedding_length(self.hparams_audio["hidden_size"])
+            self.gguf_writer.add_audio_feed_forward_length(self.hparams_audio["intermediate_size"])
+            self.gguf_writer.add_audio_head_count(self.hparams_audio["num_attention_heads"])
+            self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
+            # mimi convolves the waveform directly, it is passed around as a 1-row "mel"
+            self.gguf_writer.add_audio_num_mel_bins(1)
 
         # generation: flow-matching decoder + mimi decoder
         # the SEANet and flow net hparams are constant across the family, clip.cpp holds them
@@ -229,12 +331,57 @@ class PocketTTSMmprojModel(MmprojModel):
         # conv1d/conv1d_dw kernels must be F16, ggml_conv_1d(_dw) has no BF16 path
         if ".seanet." in new_name or new_name in ("a.downsample.conv.weight", "a.gen.wav.upsample.weight"):
             return gguf.GGMLQuantizationType.F16
+        # voices are read on the host
+        if new_name.startswith("voice."):
+            return gguf.GGMLQuantizationType.F32
         return False
 
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        # embeddings/<name>.safetensors hold the K/V state after the voice prompt (see get_state_for_audio_prompt() in the reference), turn it back into prompt embeddings
+        paths = sorted((self.dir_model / "embeddings").glob("*.safetensors"))
+        if not paths:
+            return
+        weights = _load_safetensors(self.dir_model / "model.safetensors",
+                                    ("flow_lm.transformer.layers.", "flow_lm.speaker_proj_weight", "flow_lm.bos_before_voice"))
+        n_layer = self.global_config["num_hidden_layers"]
+        for path in paths:
+            state = _load_safetensors(path, ("transformer.layers.",))
+            kv = []
+            for il in range(n_layer):
+                prefix = f"transformer.layers.{il}.self_attn/"
+                if any(prefix + k not in state for k in ("cache", "offset")):
+                    break
+                n_pos = int(state[prefix + "offset"].item())
+                # older packs have no pad entry
+                if prefix + "pad" in state and int(state[prefix + "pad"].item()) != 0:
+                    break
+                kv.append(state[prefix + "cache"][:, 0, :n_pos].flatten(2))  # [k/v, pos, n_embd]
+            if len(kv) != n_layer or any(t.shape != kv[0].shape for t in kv) or kv[0].shape[1] < 2 or \
+                    kv[0].shape[2] != self.global_config["hidden_size"]:
+                logger.warning(f"voice {path.stem}: skipped, not an unpadded voice state of this backbone")
+                continue
+            embd, drift = _voice_prompt_embd(weights, kv, self.global_config["num_attention_heads"],
+                                             self.global_config["rope_theta"])
+            # some packs ship a state made with other weights, the reference still uses it: keep up to 50% error, warn above 5%
+            if drift > 0.5:
+                logger.warning(f"voice {path.stem}: skipped, its state does not belong to these weights (error {drift:.1e})")
+                continue
+            log = logger.warning if drift > 0.05 else logger.info
+            log(f"voice {path.stem}: {embd.shape[0]} frames, prompting them again gives the shipped state within {drift:.1e}")
+            yield (f"voice.{path.stem}", embd.float())
+
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
         del bid  # the block index of the mimi transformers is parsed here, not by the base class
         T = gguf.MODEL_TENSOR
 
+        if name.startswith("voice."):
+            yield (name, data_torch)
+            return
+
+        if not self.has_audio_encoder and (name.startswith(("mimi.encoder", "mimi.downsample."))
+                                           or name == "flow_lm.speaker_proj_weight"):
+            return  # voice-prompt encoder
+
         if name in ("flow_lm.bos_emb", "flow_lm.bos_before_voice", "flow_lm.conditioner.embed.weight"):
             return  # folded into the backbone embedding table
         if name.startswith("flow_lm.transformer.") or name.startswith("flow_lm.out_norm."):

docs/backend/CANN.md

@@ -161,7 +161,7 @@ lspci -n | grep -Eo '19e5:d[0-9a-f]{3}' | cut -d: -f2
 | FP16                   | Support | Support |
 | Q8_0                   | Support | Partial |
 | Q4_0                   | Support | Partial |
-| BF16                   | Support |         |
+| BF16                   | Support |    No   |
 
 > **310P note**
 > - `Q8_0`: data transform / buffer path is implemented, and `GET_ROWS` is supported, but quantized `MUL_MAT` / `MUL_MAT_ID` are not supported.
@@ -274,7 +274,9 @@ cmake --build build --config release
 
     **Notes**:
 
-      - CANN backend only supports FP16/Q4_0/Q8_0 models currently.
+      - On 910B the CANN backend supports FP16, BF16 and Q4_0/Q8_0 models. On 310P, BF16 is not
+        supported and quantized `MUL_MAT` / `MUL_MAT_ID` are not implemented, so only FP16 models
+        are supported (see [DataType Supports](#datatype-supports)).
 
 2. **Launch inference**
 
@@ -355,3 +357,7 @@ Enable ACL graph execution during the prefill stage, default is false. This opti
 ### GGML_CANN_OPERATOR_FUSION
 
 Enable operator fusion during computation, default is false. This option fuses compatible operators (e.g., ADD + RMS_NORM) to reduce overhead and improve performance.
+
+### GGML_CANN_NO_PINNED
+
+Disable pinned (host-locked) host buffers. When it is set (any value), `aclrtMallocHost` is not called: the CANN host buffer type allocates regular CPU memory instead and the backend reports no `host_buffer` capability, so host-side staging uses pageable memory. Use it when pinned host allocation fails, or when the pinned memory is needed by other processes on the host.

docs/build-s390x.md

@@ -249,7 +249,7 @@ IBM VXE/VXE2 SIMD acceleration depends on the BLAS implementation. It is strongl
 | MXFP4      | ✅           | ❓    | ❓     |
 | Q5_0       | ✅           | ❓    | ❓     |
 | Q5_1       | ✅           | ❓    | ❓     |
-| Q8_0       | ✅           | ❓    | ❓     |
+| Q8_0       | ✅           | ✅    | ❓     |
 | Q2_K       | 🚫           | ❓    | ❓     |
 | Q3_K       | ✅           | ❓    | ❓     |
 | Q4_K       | ✅           | ❓    | ❓     |
@@ -266,11 +266,11 @@ IBM VXE/VXE2 SIMD acceleration depends on the BLAS implementation. It is strongl
 | IQ1_M      | 🚫           | ❓    | ❓     |
 | IQ4_NL     | ✅           | ❓    | ❓     |
 | IQ4_XS     | ✅           | ❓    | ❓     |
-| FP32->FP16 | 🚫           | ❓    | ❓     |
+| FP32->FP16 | ✅           | ❓    | ❓     |
 | FP16->FP32 | 🚫           | ❓    | ❓     |
 
 -   ✅ - acceleration available
 -   🚫 - acceleration unavailable, will still run using scalar implementation
 -   ❓ - acceleration unknown, please contribute if you can test it yourself
 
-Last Updated by **Aaron Teo (aaron.teo1@ibm.com)** on Sep 8, 2026.
+Last Updated by **Aaron Teo (aaron.teo1@ibm.com)** on Nov 3, 2026.

docs/build.md

@@ -481,7 +481,7 @@ cmake --build build --config Release
 Now you can load the model in conversation mode using `Vulkan`
 
 ```sh
-build/bin/Release/llama-cli -m "[PATH TO MODEL]" -ngl 100 -c 16384 -t 10 -n -2 -cnv
+build/bin/Release/llama-cli -m "[PATH TO MODEL]" -ngl 100 -c 16384 -t 10 -n -2
 ```
 
 **MSYS2**

docs/ops.md

@@ -19,7 +19,7 @@ Legend:
 |                              ADD | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             ADD1 | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                           ADD_ID | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                           ARANGE | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                           ARANGE | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           ARGMAX | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          ARGSORT | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             CEIL | ❌ | ❌ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |

docs/ops/WebGPU.csv

@@ -12326,8 +12326,16 @@
 "WebGPU: WebGPU","PAD_REFLECT_1D","type=f32,ne_a=[512,34,2,1],pad_0=10,pad_1=9","support","0","no","WebGPU"
 "WebGPU: WebGPU","PAD_REFLECT_1D","type=f32,ne_a=[3000,384,4,1],pad_0=10,pad_1=9","support","0","no","WebGPU"
 "WebGPU: WebGPU","ROLL","shift0=3,shift1=-2,shift3=1,shift4=-1","support","0","no","WebGPU"
-"WebGPU: WebGPU","ARANGE","type=f32,start=0.000000,stop=10.000000,step=1.000000","support","0","no","WebGPU"
-"WebGPU: WebGPU","ARANGE","type=f32,start=0.000000,stop=1048576.000000,step=1.000000","support","0","no","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=0.000000,stop=10.000000,step=1.000000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=3.000000,stop=3.500000,step=1.000000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-10.000000,stop=10.000000,step=0.300000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=62.750000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=63.000000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=63.250000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=254.750000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=255.000000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=255.250000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=0.000000,stop=1048576.000000,step=1.000000","support","1","yes","WebGPU"
 "WebGPU: WebGPU","TIMESTEP_EMBEDDING","type=f32,ne_a=[2,1,1,1],dim=320,max_period=10000","support","0","no","WebGPU"
 "WebGPU: WebGPU","LEAKY_RELU","type=f32,ne_a=[10,5,4,3],negative_slope=0.100000","support","0","no","WebGPU"
 "WebGPU: WebGPU","CUMSUM","type=f32,ne=[10,5,4,3]","support","1","yes","WebGPU"

ggml/include/ggml-backend.h

@@ -211,6 +211,18 @@ extern "C" {
 
     // Context management and operations for faster communication between backends, used for tensor parallelism (meta backend)
     typedef void * (*ggml_backend_comm_init_t)(ggml_backend_t * backends, size_t n_backends);
+
+    // Builds the communicator for the case where each rank is a separate process, so a process
+    // owns one backend of the communicator rather than all of them.
+    //
+    // The caller provides bootstrap_callback, which moves bytes between the ranks before the
+    // communicator exists. Each rank passes size bytes in send_buf, and recv_buf receives
+    // world_size * size bytes in rank order. The backend decides what to exchange and how many
+    // times to call it, and every rank must make the same calls in the same order.
+    // comm_init_rank blocks until all world_size ranks have joined.
+    typedef bool   (*ggml_backend_comm_bootstrap_callback)(void * user_data, const void * send_buf, void * recv_buf, size_t size);
+    typedef void * (*ggml_backend_comm_init_rank_t)(ggml_backend_t backend, int rank, int world_size, ggml_backend_comm_bootstrap_callback bootstrap_callback, void * bootstrap_callback_data);
+
     typedef void   (*ggml_backend_comm_free_t)(void * comm_ctx);
     typedef bool   (*ggml_backend_comm_allreduce_tensor_t)(void * comm_ctx, struct ggml_tensor ** tensors);
 

ggml/include/ggml.h

@@ -2516,6 +2516,24 @@ extern "C" {
             float                 max_bias,
             float                 logit_softcap);
 
+    // same as ggml_flash_attn_ext, but each q slice attends to its own list of K/V rows
+    // q:       [n_embd_k, n_batch, n_head,    ne3]
+    // k:       [n_embd_k, n_rows,  n_head_kv, 1  ] shared by all q slices
+    // v:       [n_embd_v, n_rows,  n_head_kv, 1  ]
+    // mask:    [n_kv,     n_batch, ne32,      ne3]
+    // kv_rows: [n_kv,     ne3] I32, mask column i of slice i3 uses row kv_rows[i, i3] of k and v
+    // negative rows are padding, their mask entries must be -inf
+    GGML_API struct ggml_tensor * ggml_flash_attn_ext_rows(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * q,
+            struct ggml_tensor  * k,
+            struct ggml_tensor  * v,
+            struct ggml_tensor  * mask,
+            struct ggml_tensor  * kv_rows,
+            float                 scale,
+            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),

ggml/src/ggml-backend-meta.cpp

@@ -489,6 +489,32 @@ static struct ggml_tensor * ggml_backend_meta_buffer_simple_tensor(const struct
     return it->second[index];
 }
 
+// bytes spanned by all dims except skip_dim
+static size_t ggml_backend_meta_inner_extent(const struct ggml_tensor * tensor, int skip_dim) {
+    size_t extent = 0;
+    for (int i = 0; i < GGML_MAX_DIMS; i++) {
+        if (i != skip_dim && tensor->ne[i] > 1) {
+            extent += (size_t) (tensor->ne[i] - 1) * tensor->nb[i];
+        }
+    }
+    return extent;
+}
+
+// the view dim that holds the split axis of src, or -1
+// all other dims of src and view must stay inside one slab of that axis
+static int ggml_backend_meta_view_slab_dim(const struct ggml_tensor * view, const struct ggml_tensor * src, int axis) {
+    const size_t slab = src->nb[axis];
+    if (ggml_backend_meta_inner_extent(src, axis) >= slab) {
+        return -1;
+    }
+    for (int dim = 0; dim < GGML_MAX_DIMS; dim++) {
+        if (view->ne[dim] == src->ne[axis] && view->nb[dim] == slab) {
+            return ggml_backend_meta_inner_extent(view, dim) < slab ? dim : -1;
+        }
+    }
+    return -1;
+}
+
 static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync);
 
 static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
@@ -700,9 +726,39 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             }
             GGML_ABORT("fatal error");
         }
+        // permuted view: split the view dim that spans the src split axis, other dims must be inside one element of it or outside the span
+        if (ggml_is_permuted(tensor) && !ggml_is_permuted(tensor->src[0]) && axis >= 0 && axis < GGML_MAX_DIMS-1) {
+            const size_t nb_axis = tensor->src[0]->nb[axis];
+            const size_t nb_next = tensor->src[0]->nb[axis+1];
+            for (int dim = 0; dim < GGML_MAX_DIMS; dim++) {
+                if (tensor->ne[dim] == 1 || tensor->nb[dim] < nb_axis || tensor->nb[dim]*tensor->ne[dim] != nb_next) {
+                    continue;
+                }
+                bool others_ok = true;
+                for (int i = 0; i < GGML_MAX_DIMS; i++) {
+                    if (i == dim || tensor->ne[i] == 1) {
+                        continue;
+                    }
+                    others_ok = others_ok && (tensor->nb[i]*tensor->ne[i] <= tensor->nb[dim] || tensor->nb[i] % nb_next == 0);
+                }
+                if (others_ok) {
+                    return {ggml_backend_meta_split_axis(dim), {0}, {1}, 1};
+                }
+            }
+        }
         if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED || src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) {
             return src_ss[0];
         }
+        // a view that only regroups dims inside one slab of the split axis keeps the split,
+        // the per-device strides are only rescaled above nb[axis] (e.g. the KV streams of a permuted q)
+        if (axis >= 0 && axis < GGML_MAX_DIMS && tensor->view_offs == tensor->src[0]->view_offs) {
+            const int dim = ggml_backend_meta_view_slab_dim(tensor, tensor->src[0], axis);
+            if (dim >= 0) {
+                ggml_backend_meta_split_state ret = src_ss[0];
+                ret.axis = ggml_backend_meta_split_axis(dim);
+                return ret;
+            }
+        }
         GGML_ABORT("view of permuted tensor not implemented");
         //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1};
     };
@@ -1000,8 +1056,13 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
                 split_state = handle_pad(src_ss);
             } break;
             case GGML_OP_PAD_REFLECT_1D:
-            case GGML_OP_ROLL:
-            case GGML_OP_ARANGE:
+            case GGML_OP_ROLL: {
+                split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+            } break;
+            case GGML_OP_ARANGE: {
+                // no sources, every device computes the full range
+                split_state = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
+            } break;
             case GGML_OP_TIMESTEP_EMBEDDING: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ true);
             } break;

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

@@ -2670,6 +2670,10 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
             return true;
         case GGML_OP_FLASH_ATTN_EXT:
             {
+                // K/V row lists (ggml_flash_attn_ext_rows) are not supported
+                if (op->src[5] != nullptr) {
+                    return false;
+                }
 #ifdef ASCEND_310P
                 // FA not support on 310p device
                 return false;

ggml/src/ggml-cpu/ggml-cpu-impl.h

@@ -30,7 +30,7 @@ struct ggml_compute_params {
 };
 
 
-#if defined(_MSC_VER)
+#if defined(_MSC_VER) && !defined(__clang__)
 
 #define m512bh(p) p
 #define m512i(p) p

ggml/src/ggml-cpu/ops.cpp

@@ -8844,6 +8844,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
     const ggml_tensor * v     = dst->src[2];
     const ggml_tensor * mask  = dst->src[3];
     const ggml_tensor * sinks = dst->src[4];
+    const ggml_tensor * kv_rows = dst->src[5];
 
     GGML_TENSOR_LOCALS(int64_t, neq, q,   ne)
     GGML_TENSOR_LOCALS(size_t,  nbq, q,   nb)
@@ -8951,6 +8952,8 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
         const float * pq = (const float *) ((char *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3));
         q_to_vec_dot(pq, Q_q, DK);
 
+        const int32_t * rows = kv_rows ? (const int32_t *) ((const char *) kv_rows->data + iq3*kv_rows->nb[1]) : nullptr;
+
         // online softmax / attention
         // loop over n_kv and n_head_kv
         // ref: https://arxiv.org/pdf/2112.05682.pdf
@@ -8961,9 +8964,14 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
                 continue;
             }
 
+            const int64_t ik1 = rows ? rows[ic] : ic;
+            if (ik1 < 0) {
+                continue;
+            }
+
             float s; // KQ value
 
-            const char * k_data = (const char *) k->data + ( ic*nbk1 + ik2*nbk2 + ik3*nbk3);
+            const char * k_data = (const char *) k->data + (ik1*nbk1 + ik2*nbk2 + ik3*nbk3);
             kq_vec_dot(DK, &s, 0, k_data, 0, Q_q, 0, 1);
 
             s = s*scale; // scale KQ value
@@ -8979,7 +8987,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
             float ms = 1.0f; // upon new higher max val, scale VKQ and KQ sum with this value
             float vs = 1.0f; // post-softmax KQ value, expf(s - M)
 
-            const char * v_data = ((const char *) v->data + (ic*nbv1 + iv2*nbv2 + iv3*nbv3));
+            const char * v_data = ((const char *) v->data + (ik1*nbv1 + iv2*nbv2 + iv3*nbv3));
 
             if (v->type == GGML_TYPE_F16) {
                 if (s > M) {
@@ -9146,6 +9154,10 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
     static constexpr int Q_TILE_SZ  = ggml_fa_tile_config::Q;
     static constexpr int KV_TILE_SZ = ggml_fa_tile_config::KV;
 
+    // with kv_rows the loop runs over the mask columns and reads K/V rows through kv_rows
+    const ggml_tensor * kv_rows = dst->src[5];
+    const int64_t n_kv_cols = kv_rows ? kv_rows->ne[0] : nek1;
+
     int ir = ir0;
     while (ir < ir1) {
         // q indices for the start of this tile
@@ -9153,6 +9165,10 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
         const int iq2 = (ir - iq3*neq2*neq1)/neq1;
         const int iq1 = (ir - iq3*neq2*neq1 - iq2*neq1);
 
+        // padding rows are masked, read row 0 for them
+        const int32_t * rows = kv_rows ? (const int32_t *) ((const char *) kv_rows->data + iq3*kv_rows->nb[1]) : nullptr;
+        auto kv_row = [rows](int64_t i) -> int64_t { return rows ? std::max<int32_t>(rows[i], 0) : i; };
+
         // Number of valid rows in this tile:
         // - limited by tile size (Q_TILE_SZ)
         // - limited by chunk boundary (ir1 - ir)
@@ -9212,8 +9228,8 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
         memset(K_f32, 0, DK * KV_TILE_SZ * sizeof(float));
         memset(V32,   0, KV_TILE_SZ * DV * sizeof(float));
 
-        for (int64_t ic = 0; ic < nek1; ic += KV_TILE_SZ) {
-            const int kv_tile = (int)std::min((int64_t)KV_TILE_SZ, nek1 - ic);
+        for (int64_t ic = 0; ic < n_kv_cols; ic += KV_TILE_SZ) {
+            const int kv_tile = (int)std::min((int64_t)KV_TILE_SZ, n_kv_cols - ic);
 
             // skip the tile entirely if all the masks are -inf
             if (mask) {
@@ -9240,7 +9256,7 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
             // Pack K tile transposed: K_f32[dk][kv] so KV_TILE is contiguous (SIMD dim)
             // Zero-pad the last tile so the GEMM always operates on KV_TILE_SZ columns
             for (int tk = 0; tk < kv_tile; tk++) {
-                const char * k_data = (const char *)k->data + (ic + tk)*nbk1 + ik2*nbk2 + ik3*nbk3;
+                const char * k_data = (const char *)k->data + kv_row(ic + tk)*nbk1 + ik2*nbk2 + ik3*nbk3;
                 if (kv_type == GGML_TYPE_F16) {
                     const ggml_fp16_t * k_f16 = (const ggml_fp16_t *)k_data;
                     for (int64_t dk = 0; dk < DK; dk++) {
@@ -9305,7 +9321,7 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
             // V accumulation: VKQ32 += softmax(KQ) * V
             // Pack V tile to contiguous F32, zero-padded
             for (int tk = 0; tk < kv_tile; tk++) {
-                const char * v_data = (const char *)v->data + (ic + tk)*nbv1 + iv2*nbv2 + iv3*nbv3;
+                const char * v_data = (const char *)v->data + kv_row(ic + tk)*nbv1 + iv2*nbv2 + iv3*nbv3;
                 if (kv_type == GGML_TYPE_F16) {
                     ggml_cpu_fp16_to_fp32((const ggml_fp16_t *)v_data, V32 + tk * DV, DV);
                 } else {
@@ -9477,8 +9493,12 @@ static void ggml_compute_forward_flash_attn_ext_f16(
     // When use_ref is set, force the vec-only reference implementation (no tiling, no KV-chunking)
     const bool use_ref = params->use_ref;
 
+    // with kv_rows the loops run over the mask columns, the split KV path does not support it
+    const ggml_tensor * kv_rows = dst->src[5];
+    const int64_t n_kv_cols = kv_rows ? kv_rows->ne[0] : nek1;
+
     const bool kv_is_f32_or_f16 = (k->type == GGML_TYPE_F32 || k->type == GGML_TYPE_F16);
-    const bool use_split_kv_path = !use_ref && (neq1 == 1 && neq3 == 1) && kv_is_f32_or_f16 && (k->type == v->type) && q->type == GGML_TYPE_F32 && nek1 >= 512;
+    const bool use_split_kv_path = !use_ref && !kv_rows && (neq1 == 1 && neq3 == 1) && kv_is_f32_or_f16 && (k->type == v->type) && q->type == GGML_TYPE_F32 && nek1 >= 512;
 
     if (use_split_kv_path) {
         const int64_t chunk_size = (nek1 + nth - 1) / nth;
@@ -9557,7 +9577,7 @@ static void ggml_compute_forward_flash_attn_ext_f16(
             if (use_tiled) {
                 ggml_compute_forward_flash_attn_ext_tiled(params, dst, ir0, ir1);
             } else {
-                ggml_compute_forward_flash_attn_ext_f16_one_chunk(params, dst, ir0, ir1, 0, nek1, nullptr, 0);
+                ggml_compute_forward_flash_attn_ext_f16_one_chunk(params, dst, ir0, ir1, 0, n_kv_cols, nullptr, 0);
             }
 
             current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1);

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

@@ -1046,6 +1046,10 @@ class tensor_traits_common : public tensor_traits_base {
                         return true;
                 }
             case GGML_OP_FLASH_ATTN_EXT:
+                // K/V row lists are handled by the generic CPU path
+                if (op->src[5] != nullptr) {
+                    return false;
+                }
                 forward_flash_attn_ext_f16(params, op);
                 return true;
             case GGML_OP_CONT:

ggml/src/ggml-cuda/batch-get.cu

@@ -0,0 +1,198 @@
+#include "batch-get.cuh"
+#include "common.cuh"
+#include "ggml-cpp.h"
+
+#include <climits>
+
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+struct ggml_cuda_batch_get_descriptor {
+    const unsigned char * src;
+    unsigned char * dst;
+    size_t size;
+};
+
+static __global__ void ggml_cuda_batch_get_kernel(const ggml_cuda_batch_get_descriptor * copies, size_t count) {
+    // Blocks along y cooperate on each descriptor.
+    const size_t index = blockIdx.x;
+    if (index >= count) {
+        return;
+    }
+    const size_t first = (size_t) blockIdx.y * blockDim.x + threadIdx.x;
+    const size_t stride = (size_t) gridDim.y * blockDim.x;
+    const auto copy = copies[index];
+    // Use 16-, 8-, 4-, or 2-byte copies when both pointers meet that alignment; threads copy separate chunks.
+    const uintptr_t alignment = (uintptr_t) copy.src | (uintptr_t) copy.dst;
+    size_t tail = 0;
+    if (alignment % 16 == 0) {
+        for (size_t i = first; i < copy.size / 16; i += stride) {
+            ggml_cuda_memcpy_1<16>(copy.dst + 16*i, copy.src + 16*i);
+        }
+        tail = copy.size / 16 * 16;
+    } else if (alignment % 8 == 0) {
+        for (size_t i = first; i < copy.size / 8; i += stride) {
+            ggml_cuda_memcpy_1<8>(copy.dst + 8*i, copy.src + 8*i);
+        }
+        tail = copy.size / 8 * 8;
+    } else if (alignment % 4 == 0) {
+        for (size_t i = first; i < copy.size / 4; i += stride) {
+            ggml_cuda_memcpy_1<4>(copy.dst + 4*i, copy.src + 4*i);
+        }
+        tail = copy.size / 4 * 4;
+    } else if (alignment % 2 == 0) {
+        for (size_t i = first; i < copy.size / 2; i += stride) {
+            ggml_cuda_memcpy_1<2>(copy.dst + 2*i, copy.src + 2*i);
+        }
+        tail = copy.size / 2 * 2;
+    }
+    // Copy trailing bytes, or the entire range when no wider alignment is available.
+    for (size_t i = tail + first; i < copy.size; i += stride) {
+        ggml_cuda_memcpy_1<1>(copy.dst + i, copy.src + i);
+    }
+}
+
+struct ggml_cuda_batch_get_mapping {
+    void * host = nullptr;
+    unsigned char * device = nullptr;
+
+    ggml_cuda_batch_get_mapping(void * host, size_t size) {
+        const cudaError_t err = cudaHostRegister(host, size, cudaHostRegisterMapped | cudaHostRegisterPortable);
+        if (err != cudaSuccess) {
+            (void) cudaGetLastError();
+            GGML_LOG_DEBUG("%s: failed to register %.2f MiB of pinned memory: %s\n", __func__,
+                           size / 1024.0 / 1024.0, cudaGetErrorString(err));
+            return;
+        }
+        this->host = host;
+        if (cudaHostGetDevicePointer((void **) &device, host, 0) != cudaSuccess) {
+            (void) cudaGetLastError();
+            device = nullptr;
+            return;
+        }
+    }
+
+    ggml_cuda_batch_get_mapping(const ggml_cuda_batch_get_mapping &) = delete;
+    ggml_cuda_batch_get_mapping & operator=(const ggml_cuda_batch_get_mapping &) = delete;
+
+    ~ggml_cuda_batch_get_mapping() {
+        if (host) {
+            // Release only the registration acquired by this object.
+            CUDA_CHECK(cudaHostUnregister(host));
+        }
+    }
+};
+#endif
+
+bool ggml_cuda_get_tensor_n(int device,
+        const ggml_tensor * const * tensors, void * const * data, const size_t * offsets,
+        const size_t * sizes, size_t n_tensors) {
+#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
+    GGML_UNUSED(device);
+    GGML_UNUSED(tensors);
+    GGML_UNUSED(data);
+    GGML_UNUSED(offsets);
+    GGML_UNUSED(sizes);
+    GGML_UNUSED(n_tensors);
+    return false;
+#else
+    // Provisional dispatch values based on experiments.
+    constexpr size_t min_batch_reads = 128;
+    constexpr size_t max_gather_read_bytes = 128 * 1024;
+    // Collect non-empty CPU destination ranges.
+    std::vector<std::pair<uintptr_t, uintptr_t>> ranges;
+    ranges.reserve(n_tensors);
+    size_t payload_bytes = 0;
+    // Number of non-empty reads whose size is at most max_gather_read_bytes.
+    size_t n_reads_small = 0;
+    for (size_t i = 0; i < n_tensors; ++i) {
+        if (sizes[i] == 0) {
+            continue;
+        }
+        const uintptr_t address = (uintptr_t) data[i];
+        GGML_ASSERT(address <= UINTPTR_MAX - sizes[i]);
+        GGML_ASSERT(payload_bytes <= SIZE_MAX - sizes[i]);
+        payload_bytes += sizes[i];
+        n_reads_small += sizes[i] <= max_gather_read_bytes;
+        ranges.emplace_back(address, address + sizes[i]);
+    }
+    if (ranges.empty()) {
+        return true;
+    }
+    // Total number of non-empty reads.
+    const size_t n_reads = ranges.size();
+    if (n_reads > SIZE_MAX / sizeof(ggml_cuda_batch_get_descriptor) || n_reads > INT_MAX) {
+        return false;
+    }
+    if (n_reads < 2 || (n_reads < min_batch_reads && n_reads_small == n_reads)) {
+        return false;
+    }
+    std::sort(ranges.begin(), ranges.end());
+    for (size_t i = 1; i < ranges.size(); ++i) {
+        // Preserve the ordered copy path when CPU destinations overlap.
+        if (ranges[i].first < ranges[i - 1].second) {
+            return false;
+        }
+    }
+    // Find the CPU span to register, including gaps between destinations.
+    const uintptr_t begin = ranges.front().first;
+    const uintptr_t end = ranges.back().second;
+    if (end - begin - payload_bytes > payload_bytes) {
+        return false;
+    }
+    // Register and map the existing CPU destination so the kernel can write to it directly.
+    ggml_cuda_batch_get_mapping mapping((void *) begin, end - begin);
+    if (!mapping.host) {
+        return false;
+    }
+    // Reads selected for gather; zero means all reads use DMA.
+    size_t n_reads_gather = mapping.device && n_reads_small >= min_batch_reads ? n_reads_small : 0;
+    // Keep descriptor storage alive until the final stream synchronization.
+    constexpr size_t max_batch_descriptors = 128 * 1024;
+    const size_t bytes = std::min(n_reads_gather, max_batch_descriptors) * sizeof(ggml_cuda_batch_get_descriptor);
+    ggml_backend_buffer_ptr descriptors(n_reads_gather ? ggml_backend_buft_alloc_buffer(ggml_backend_cuda_buffer_type(device), bytes) : nullptr);
+    if (!descriptors) {
+        n_reads_gather = 0;
+    }
+    const auto host_type = ggml_backend_cuda_host_buffer_type();
+    ggml_backend_buffer_ptr host_descriptors(n_reads_gather ? ggml_backend_buft_alloc_buffer(host_type, bytes) : nullptr);
+    if (!host_descriptors || ggml_backend_buffer_get_type(host_descriptors.get()) != host_type) {
+        n_reads_gather = 0;
+    }
+    if (n_reads_gather) {
+        auto * copies = static_cast<ggml_cuda_batch_get_descriptor *>(ggml_backend_buffer_get_base(host_descriptors.get()));
+        auto * device_copies = static_cast<ggml_cuda_batch_get_descriptor *>(ggml_backend_buffer_get_base(descriptors.get()));
+        const int threads = 256;
+        const size_t nsm = ggml_cuda_info().devices[device].nsm;
+        size_t i = 0;
+        for (size_t first = 0; first < n_reads_gather; first += max_batch_descriptors) {
+            const size_t n_batch = std::min(n_reads_gather - first, max_batch_descriptors);
+            size_t count = 0;
+            for (; i < n_tensors && count < n_batch; ++i) {
+                if (sizes[i] == 0 || sizes[i] > max_gather_read_bytes) {
+                    continue;
+                }
+                copies[count++] = {(const unsigned char *) tensors[i]->data + offsets[i],
+                        mapping.device + ((uintptr_t) data[i] - begin), sizes[i]};
+            }
+            const size_t blocks_per_copy = std::min<size_t>(UINT16_MAX, std::max<size_t>(1, (nsm + n_batch - 1) / n_batch));
+            const dim3 blocks((unsigned int) n_batch, (unsigned int) blocks_per_copy);
+            // Stream ordering keeps each kernel ahead of the next descriptor overwrite.
+            CUDA_CHECK(cudaMemcpyAsync(device_copies, copies, n_batch * sizeof(copies[0]), cudaMemcpyHostToDevice, cudaStreamPerThread));
+            ggml_cuda_batch_get_kernel<<<blocks, threads, 0, cudaStreamPerThread>>>(device_copies, n_batch);
+            CUDA_CHECK(cudaGetLastError());
+            if (first + n_batch < n_reads_gather) {
+                // Wait before refilling the pinned host descriptor buffer.
+                CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+            }
+        }
+    }
+    // Keep the destination pinned for DMA too; overlapping destinations already declined above.
+    for (size_t i = 0; i < n_tensors; ++i) {
+        if (sizes[i] && (!n_reads_gather || sizes[i] > max_gather_read_bytes)) {
+            CUDA_CHECK(cudaMemcpyAsync(data[i], (const char *) tensors[i]->data + offsets[i], sizes[i], cudaMemcpyDeviceToHost, cudaStreamPerThread));
+        }
+    }
+    // Wait for GPU writes before releasing the CPU mapping and descriptor buffers.
+    CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+    return true;
+#endif
+}

ggml/src/ggml-cuda/batch-get.cuh

@@ -0,0 +1,7 @@
+#pragma once
+
+#include "ggml-backend.h"
+
+bool ggml_cuda_get_tensor_n(int device,
+        const ggml_tensor * const * tensors, void * const * data, const size_t * offsets,
+        const size_t * sizes, size_t n_tensors);

ggml/src/ggml-cuda/common.cuh

@@ -1469,16 +1469,21 @@ struct ggml_backend_cuda_context {
     // when the computation is split across CPU/GPU (e.g., with --n-cpu-moe)
     std::unordered_map<const void *, std::unique_ptr<ggml_cuda_graph>> cuda_graphs;
 
-    int64_t last_graph_eviction_sweep = 0;
+    int64_t next_graph_eviction_time = 0;
+    int64_t last_graph_use_time = 0;
 
-    ggml_cuda_graph * cuda_graph(const void * first_node_ptr) {
-        const int64_t time_now = ggml_time_us();
+    ggml_cuda_graph * cuda_graph(const void * first_node_ptr, int64_t time_now = ggml_time_us()) {
+        // Give all cached graphs 10s to be reused after an idle period, without updating each graph.
+        if (!cuda_graphs.empty() && time_now - last_graph_use_time >= 10'000'000) {
+            next_graph_eviction_time = time_now + 10'000'000;
+        }
+        last_graph_use_time = time_now;
 
-        // sweep every 5s, evicting cuda graphs unused for >=10s
-        if (time_now - last_graph_eviction_sweep >= 5'000'000) {
-            last_graph_eviction_sweep = time_now;
+        // Sweep every 5s during active use, evicting graphs unused for at least 10s.
+        if (time_now >= next_graph_eviction_time) {
+            next_graph_eviction_time = time_now + 5'000'000;
             for (auto it = cuda_graphs.begin(); it != cuda_graphs.end(); ) {
-                if (time_now - it->second->last_used_time >= 10'000'000) {
+                if (it->first != first_node_ptr && time_now - it->second->last_used_time >= 10'000'000) {
                     it = cuda_graphs.erase(it);
                 } else {
                     ++it;

ggml/src/ggml-cuda/fattn-common.cuh

@@ -25,6 +25,7 @@ typedef void (* fattn_kernel_t)(
         const char * __restrict__ mask,
         const char * __restrict__ sinks,
         const int  * __restrict__ KV_max,
+        const int  * __restrict__ KV_rows,
         float      * __restrict__ dst,
         float2     * __restrict__ dst_meta,
         const float scale,
@@ -991,8 +992,13 @@ void launch_fattn(
 
     const bool V_is_K_view = V->view_src && (V->view_src == K || (V->view_src == K->view_src && V->view_offs == K->view_offs));
 
-    const ggml_tensor * mask  = dst->src[3];
-    const ggml_tensor * sinks = dst->src[4];
+    const ggml_tensor * mask    = dst->src[3];
+    const ggml_tensor * sinks   = dst->src[4];
+    const ggml_tensor * kv_rows = dst->src[5];
+
+    // with kv_rows the kernel iterates over the mask columns and reads K/V rows through kv_rows
+    const int64_t n_kv_cols = kv_rows ? kv_rows->ne[0] : K->ne[1];
+    GGML_ASSERT(!kv_rows || !use_sparse);
 
     ggml_tensor * KQV = dst;
 
@@ -1113,8 +1119,8 @@ void launch_fattn(
 
     // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped.
     // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or
-    //     multiple sequences of possibly different lengths.
-    const bool scan_mask = !use_sparse && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1);
+    //     multiple sequences of possibly different lengths, or kv_rows lists that end with padding.
+    const bool scan_mask = !use_sparse && mask && n_kv_cols % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1 || kv_rows);
     if (scan_mask) {
         const int64_t s31 = mask->nb[1] / sizeof(half2);
         const int64_t s33 = mask->nb[3] / sizeof(half2);
@@ -1123,7 +1129,7 @@ void launch_fattn(
         const dim3 block_dim_KV_max(FATTN_KQ_STRIDE/2, 1, 1);
 
         const int ne_KV_max = blocks_num_KV_max.x*blocks_num_KV_max.y;
-        const int iter_k = K->ne[1] / FATTN_KQ_STRIDE;
+        const int iter_k = n_kv_cols / FATTN_KQ_STRIDE;
 
         KV_max.alloc(ne_KV_max);
         ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks_num_KV_max, block_dim_KV_max, 0, main_stream);
@@ -1132,13 +1138,19 @@ void launch_fattn(
         CUDA_CHECK(cudaGetLastError());
     }
 
+    // K/V are shared by all slices when read through kv_rows
+    if (kv_rows) {
+        nb13 = 0;
+        nb23 = 0;
+    }
+
     const dim3 block_dim(warp_size, nwarps, 1);
     int max_blocks_per_sm = 1; // Max. number of active blocks limited by occupancy.
     CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&max_blocks_per_sm, fattn_kernel, block_dim.x * block_dim.y * block_dim.z, nbytes_shared));
     GGML_ASSERT(max_blocks_per_sm > 0);
     int parallel_blocks = max_blocks_per_sm;
 
-    const int64_t n_kv = use_sparse ? n_kv_max : K->ne[1];
+    const int64_t n_kv = use_sparse ? n_kv_max : n_kv_cols;
     const int ntiles_KV = (n_kv + nbatch_fa - 1) / nbatch_fa; // Max. number of parallel blocks limited by KV cache length.
 
     dim3 blocks_num;
@@ -1255,6 +1267,7 @@ void launch_fattn(
         mask ? ((const char *) mask->data) : nullptr,
         sinks ? ((const char *) sinks->data) : nullptr,
         KV_max.ptr,
+        kv_rows ? (const int *) kv_rows->data : nullptr,
         !stream_k && parallel_blocks > 1 ? dst_tmp.ptr : (float *) KQV->data, dst_tmp_meta.ptr,
         scale, max_bias, m0, m1, n_head_log2, logit_softcap,
         Q->ne[0], ne01,     Q->ne[2], Q->ne[3], Q->nb[1], Q->nb[2], Q->nb[3],

ggml/src/ggml-cuda/fattn-mma-f16.cuh

@@ -420,6 +420,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile(
                     // padded slots gather row 0, the -inf mask removes their contribution
                     const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : 0;
                     i_KV = index >= 0 ? index : 0;
+                } else if (indices) {
+                    // K/V rows given per mask column, negative rows are padding
+                    i_KV = max(indices[k_VKQ_0 + i], 0);
                 } else {
                     i_KV = k_VKQ_0 + i;
                 }
@@ -467,8 +470,11 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile(
                     if constexpr (use_sparse) {
                         const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : -1;
                         src = index >= 0 ? KV + int64_t(index)*stride_KV + k*h2_per_chunk : zero;
+                    } else if (!oob_check || i < i_sup) {
+                        const int64_t i_KV = indices ? max(indices[k_VKQ_0 + i], 0) : k_VKQ_0 + i;
+                        src = KV + i_KV*stride_KV + k*h2_per_chunk;
                     } else {
-                        src = !oob_check || i < i_sup ? KV + int64_t(k_VKQ_0 + i)*stride_KV + k*h2_per_chunk : zero;
+                        src = zero;
                     }
                     ggml_cuda_memcpy_1<16>(swizzle<stride_tile>(tile_KV, i*stride_tile + k*h2_per_chunk, i), src);
                 }
@@ -634,7 +640,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
         cp_async_wait_all();
         __syncthreads();
         flash_attn_ext_f16_load_tile<stride_tile_V, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
-            (V_h2, tile_V, nbatch_V2, stride_V, k_VKQ_0, k_VKQ_sup, nullptr);
+            (V_h2, tile_V, nbatch_V2, stride_V, k_VKQ_0, k_VKQ_sup, indices);
     } else {
         // the sparse mask values are gathered per element, always load them synchronously
         constexpr bool use_cp_async = nstages == 1 && !use_sparse;
@@ -991,7 +997,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
                     (mask_h, tile_mask, stride_mask, k_VKQ_0 + nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr);
             }
             flash_attn_ext_f16_load_tile<stride_tile_K, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
-                (K_h2, tile_K, nbatch_K2, stride_K, k_VKQ_0 + nbatch_fa, k_VKQ_sup, nullptr);
+                (K_h2, tile_K, nbatch_K2, stride_K, k_VKQ_0 + nbatch_fa, k_VKQ_sup, indices);
         }
     }
 
@@ -1343,7 +1349,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
                 (mask_h, tile_mask, stride_mask, kb0*nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr);
         }
         flash_attn_ext_f16_load_tile<stride_tile_K, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
-            (K_h2, tile_K, nbatch_K2, stride_K, kb0*nbatch_fa, k_VKQ_sup, nullptr);
+            (K_h2, tile_K, nbatch_K2, stride_K, kb0*nbatch_fa, k_VKQ_sup, indices);
     }
 
     // kb0_start is always < kb0_stop so the last iter can be executed unconditionally.
@@ -1802,6 +1808,7 @@ static __global__ void flash_attn_ext_f16(
         const char * mask_ptr,
         const char * sinks_ptr,
         const int  * KV_max_ptr,
+        const int  * KV_rows_ptr,
         float      * dst_ptr,
         float2     * dst_meta_ptr,
         const float scale,
@@ -1924,7 +1931,8 @@ static __global__ void flash_attn_ext_f16(
 
         const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV);
         const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr;
-        const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*iter_j + jt)*ne11 : nullptr;
+        const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*iter_j + jt)*ne11 :
+                                  KV_rows_ptr ? KV_rows_ptr + int64_t(sequence)*ne11 : nullptr;
 
         const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f;
 
@@ -1973,7 +1981,8 @@ static __global__ void flash_attn_ext_f16(
 
     const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV);
     const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr;
-    const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*iter_j + jt)*ne11 : nullptr;
+    const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*iter_j + jt)*ne11 :
+                              KV_rows_ptr ? KV_rows_ptr + int64_t(sequence)*ne11 : nullptr;
 
     const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f;
 
@@ -1989,7 +1998,7 @@ static __global__ void flash_attn_ext_f16(
         (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap,
          ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop);
 #else
-    GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale,
+    GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, KV_rows_ptr, dst_ptr, dst_meta_ptr, scale,
         max_bias, m0, m1, n_head_log2, logit_softcap,
         ne00, ne01, ne02, ne03,
               nb01, nb02, nb03,

ggml/src/ggml-cuda/fattn-tile.cuh

@@ -797,6 +797,7 @@ static __global__ void flash_attn_tile(
         const char * mask_ptr,
         const char * sinks_ptr,
         const int  * KV_max_ptr,
+        const int  * KV_rows_ptr,
         float      * dst_ptr,
         float2     * dst_meta_ptr,
         const float scale,
@@ -821,6 +822,7 @@ static __global__ void flash_attn_tile(
     const int  * GGML_CUDA_RESTRICT KV_max   = KV_max_ptr;
     float      * GGML_CUDA_RESTRICT dst      = dst_ptr;
     float2     * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr;
+    GGML_UNUSED(KV_rows_ptr);
 
     // Skip unused kernel variants for faster compilation:
 
@@ -1132,7 +1134,7 @@ static __global__ void flash_attn_tile(
         }
     }
 #else
-    GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale,
+    GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, KV_rows_ptr, dst_ptr, dst_meta_ptr, scale,
         max_bias, m0, m1, n_head_log2, logit_softcap,
         ne00, ne01, ne02, ne03,
               nb01, nb02, nb03,

ggml/src/ggml-cuda/fattn-vec.cuh

@@ -25,6 +25,7 @@ static __global__ void flash_attn_ext_vec(
         const char * mask_ptr,
         const char * sinks_ptr,
         const int  * KV_max_ptr,
+        const int  * KV_rows_ptr,
         float      * dst_ptr,
         float2     * dst_meta_ptr,
         const float scale,
@@ -50,6 +51,7 @@ static __global__ void flash_attn_ext_vec(
     const int  * GGML_CUDA_RESTRICT KV_max   = KV_max_ptr;
     float      * GGML_CUDA_RESTRICT dst      = dst_ptr;
     float2     * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr;
+    GGML_UNUSED(KV_rows_ptr);
 
     // Skip unused kernel variants for faster compilation:
     if (use_logit_softcap && !(D == 128 || D == 256)) {
@@ -514,7 +516,7 @@ static __global__ void flash_attn_ext_vec(
         dst_meta[((sequence*int(ne01.z) + ic0 + tid)*ne02 + head)*gridDim.y + blockIdx.y] = make_float2(KQ_max[tid], KQ_sum[tid]);
     }
 #else
-    GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale,
+    GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, KV_rows_ptr, dst_ptr, dst_meta_ptr, scale,
         max_bias, m0, m1, n_head_log2, logit_softcap,
         ne00, ne01, ne02, ne03,
               nb01, nb02, nb03,

ggml/src/ggml-cuda/fattn.cu

@@ -161,7 +161,7 @@ bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(const int cc, const ggml_
     // the dense kernel handles up to 64/ncols2 queries per K/V pass, the single-query gather has to beat that
     const int64_t n_gather = (ncols1 == 1 ? std::min<int64_t>(Q->ne[1], 64/ncols2) : ncols1) * (int64_t) n_kv_max;
 
-    return ((GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc)) || amd_mfma_available(cc)) &&
+    return ((GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc)) || amd_mfma_available(cc)) && dst->src[5] == nullptr &&
         mask != nullptr && n_kv_max > 0 && max_bias == 0.0f && logit_softcap == 0.0f &&
         mask->ne[0] == K->ne[1] && mask->ne[1] >= Q->ne[1] && mask->ne[2] == 1 &&
         K->ne[1] >= std::max<int64_t>(4096, GGML_CUDA_CC_IS_NVIDIA(cc) ? 2*n_gather : 2LL*n_kv_max);
@@ -216,12 +216,15 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2(ggml_backend_cuda_con
     const ggml_tensor * V    = dst->src[2];
     const ggml_tensor * mask = dst->src[3];
 
+    // with kv_rows the kernel iterates over the mask columns
+    const int64_t n_kv = dst->src[5] ? dst->src[5]->ne[0] : K->ne[1];
+
     float max_bias = 0.0f;
     memcpy(&max_bias, (const float *) KQV->op_params + 1, sizeof(float));
 
     // Edge cases like no mask, ALiBi, unpadded K/V, or misaligned addresses for large data transfers
     //     are put into the template specialization without GQA optimizations.
-    bool use_gqa_opt = mask && max_bias == 0.0f && K->ne[1] % FATTN_KQ_STRIDE == 0;
+    bool use_gqa_opt = mask && max_bias == 0.0f && n_kv % FATTN_KQ_STRIDE == 0;
     for (const ggml_tensor * t : {Q, K, V, mask}) {
         if (t == nullptr || ggml_is_quantized(t->type)) {
             continue;
@@ -567,6 +570,10 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
     const ggml_tensor * K     = dst->src[1];
     const ggml_tensor * V     = dst->src[2];
     const ggml_tensor * mask  = dst->src[3];
+    const ggml_tensor * kv_rows = dst->src[5];
+
+    // with kv_rows the kernel iterates over the mask columns
+    const int64_t n_kv = kv_rows ? kv_rows->ne[0] : K->ne[1];
 
     const int gqa_ratio = Q->ne[2] / K->ne[2];
     GGML_ASSERT(Q->ne[2] % K->ne[2] == 0);
@@ -576,7 +583,7 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
 
     // The effective batch size for the kernel can be increased by gqa_ratio.
     // The kernel versions without this optimization are also used for ALiBi, if there is no mask, or if the KV cache is not padded,
-    bool gqa_opt_applies = gqa_ratio >= 2 && mask && max_bias == 0.0f && K->ne[1] % FATTN_KQ_STRIDE == 0;
+    bool gqa_opt_applies = gqa_ratio >= 2 && mask && max_bias == 0.0f && n_kv % FATTN_KQ_STRIDE == 0;
     for (const ggml_tensor * t : {Q, K, V, mask}) {
         if (t == nullptr || ggml_is_quantized(t->type)) {
             continue;
@@ -648,6 +655,12 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
         return BEST_FATTN_KERNEL_NONE;
     }
 
+    // only the MMA kernel reads K/V through kv_rows
+    if (kv_rows) {
+        return GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && mask && Q->ne[0] != 40 && Q->ne[0] != 72 ?
+            BEST_FATTN_KERNEL_MMA_F16 : BEST_FATTN_KERNEL_NONE;
+    }
+
     // For small batch sizes the vector kernel may be preferable over the kernels optimized for large batch sizes:
     // 192 satisfies % 64 == 0 but has no vec instance (DKQ != DV); force it onto the MMA path.
     const bool can_use_vector_kernel = Q->ne[0] <= 256 && Q->ne[0] % 64 == 0 && Q->ne[0] != 192 && K->ne[1] % FATTN_KQ_STRIDE == 0;

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

@@ -9,6 +9,7 @@
 #include "ggml-cuda/arange.cuh"
 #include "ggml-cuda/argmax.cuh"
 #include "ggml-cuda/argsort.cuh"
+#include "ggml-cuda/batch-get.cuh"
 #include "ggml-cuda/binbcast.cuh"
 #include "ggml-cuda/clamp.cuh"
 #include "ggml-cuda/col2im-1d.cuh"
@@ -886,6 +887,11 @@ static void ggml_backend_cuda_buffer_get_tensor_n(ggml_backend_buffer_t buffer,
 
     ggml_cuda_set_device(ctx->device);
 
+    if (!ggml_cuda_pinned_memory_disabled() &&
+            ggml_cuda_get_tensor_n(ctx->device, tensors, data, offsets, sizes, n_tensors)) {
+        return;
+    }
+
     for (size_t i = 0; i < n_tensors; i++) {
         if (sizes[i] == 0) {
             continue;
@@ -1117,6 +1123,18 @@ struct ggml_backend_cuda_comm_context {
 
     ggml_cuda_ar_pipeline *     ar_pipeline = nullptr;
 
+    // Number of ranks in the communicator when each rank is a separate process. Zero unless
+    // comm_init_rank created this context.
+    int                         world_size = 0;
+
+    // Number of ranks taking part in the reduction. For a communicator from comm_init_rank this is
+    // world_size. Otherwise it is the number of backends this process holds. The NCCL reduction
+    // uses it to choose the tensor size below which it reduces in F32 rather than converting to
+    // BF16 first.
+    size_t n_ranks() const {
+        return world_size > 0 ? (size_t) world_size : backends.size();
+    }
+
 #ifdef GGML_USE_NCCL
     std::vector<ncclComm_t>     comms;
 #endif // GGML_USE_NCCL
@@ -1144,6 +1162,7 @@ static bool ggml_backend_cuda_comm_allreduce_nccl(
     }
 
     const size_t n_backends = comm_ctx->backends.size();
+    const size_t n_ranks    = comm_ctx->n_ranks();
 
     for (size_t i = 0; i < n_backends; ++i) {
         GGML_ASSERT(tensors[i] != nullptr);
@@ -1153,7 +1172,7 @@ static bool ggml_backend_cuda_comm_allreduce_nccl(
 
     // For small tensors, simply reduce them as FP32.
     // The following heuristic for how "small" a tensor should be is based on RTX 4090s connected via 16x PCIe 4.0.
-    if ((n_backends <= 2 && ne < 32768) || (n_backends == 3 && ne < 131072) || (n_backends >= 4 && ne < 262144)) {
+    if ((n_ranks <= 2 && ne < 32768) || (n_ranks == 3 && ne < 131072) || (n_ranks >= 4 && ne < 262144)) {
         for (size_t i = 0; i < n_backends; ++i) {
             if ((tensors[i]->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
                 ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context;
@@ -1384,6 +1403,90 @@ static void * ggml_backend_cuda_comm_init(ggml_backend_t * backends, size_t n_ba
     return ret;
 }
 
+#ifdef GGML_USE_NCCL
+// Builds the communicator with ncclCommInitRank, for the case where each rank is its own
+// process. Blocks until all world_size ranks have joined.
+static void * ggml_backend_cuda_comm_init_rank(ggml_backend_t backend, int rank, int world_size,
+                                               ggml_backend_comm_bootstrap_callback bootstrap_callback, void * bootstrap_callback_data) {
+    if (world_size < 1 || rank < 0 || rank >= world_size) {
+        GGML_LOG_ERROR("%s: invalid rank %d for world size %d\n", __func__, rank, world_size);
+        return nullptr;
+    }
+
+    // Every rank sends whether its local checks passed, and rank 0 also sends the NCCL unique id.
+    // So all ranks either join together or fail together, and no rank waits in ncclCommInitRank alone.
+    struct bootstrap_msg {
+        uint8_t      ok;
+        ncclUniqueId id;
+    };
+    bootstrap_msg local;
+    memset(&local, 0, sizeof(local));
+    local.ok = 1;
+
+    if (!ggml_backend_is_cuda(backend)) {
+        GGML_LOG_ERROR("%s: backend is not a CUDA backend\n", __func__);
+        local.ok = 0;
+    } else {
+        // NCCL requires one distinct physical GPU per rank, see ggml_backend_cuda_comm_init_nccl.
+        const ggml_cuda_device_info & info = ggml_cuda_info();
+        if (info.device_count > info.physical_device_count) {
+            GGML_LOG_ERROR("%s: CUDA virtual devices are in use, which NCCL does not support\n", __func__);
+            local.ok = 0;
+        }
+    }
+    if (local.ok && rank == 0) {
+        const ncclResult_t rc = ncclGetUniqueId(&local.id);
+        if (rc != ncclSuccess) {
+            GGML_LOG_ERROR("%s: ncclGetUniqueId failed (%s)\n", __func__, ncclGetErrorString(rc));
+            local.ok = 0;
+        }
+    }
+
+    std::vector<bootstrap_msg> all(world_size);
+    if (!bootstrap_callback(bootstrap_callback_data, &local, all.data(), sizeof(local))) {
+        GGML_LOG_ERROR("%s: bootstrap exchange failed on rank %d\n", __func__, rank);
+        return nullptr;
+    }
+    for (int i = 0; i < world_size; i++) {
+        if (!all[i].ok) {
+            GGML_LOG_ERROR("%s: rank %d failed its setup checks\n", __func__, i);
+            return nullptr;
+        }
+    }
+
+    ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context;
+
+    auto * ret = new ggml_backend_cuda_comm_context;
+    ret->backends.push_back(backend);
+    ret->dev_ids.push_back(cuda_ctx->device);
+    ret->world_size = world_size;
+    ret->comms.resize(1);
+    ret->try_allreduce = ggml_backend_cuda_comm_try_allreduce_nccl;
+
+    // On DGX Spark, NCCL picks 64 channels by default, which is slower and uses more memory than 8.
+    // NCCL_MAX_CTAS still overrides this.
+    ncclConfig_t config = NCCL_CONFIG_INITIALIZER;
+    const int cc = ggml_cuda_info().devices[cuda_ctx->device].cc;
+    if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_DGX_SPARK) {
+        config.maxCTAs = 8;
+    }
+
+    ggml_cuda_set_device(cuda_ctx->device);
+    const ncclResult_t rc = ncclCommInitRankConfig(&ret->comms[0], world_size, all[0].id, rank, &config);
+    if (rc != ncclSuccess) {
+        GGML_LOG_ERROR("%s: ncclCommInitRankConfig failed for rank %d/%d (%s)\n",
+                       __func__, rank, world_size, ncclGetErrorString(rc));
+        ret->comms.clear();
+        delete ret;
+        return nullptr;
+    }
+
+    GGML_LOG_INFO("%s: joined NCCL communicator as rank %d/%d on device %d\n",
+                  __func__, rank, world_size, cuda_ctx->device);
+    return ret;
+}
+#endif // GGML_USE_NCCL
+
 // Top-level dispatch -- calls the function pointer chosen by comm_init.
 // Returns false to let the meta-backend's butterfly run.
 static bool ggml_backend_cuda_comm_allreduce_tensor(void * comm_ctx_v, struct ggml_tensor ** tensors) {
@@ -6384,6 +6487,11 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con
     if (strcmp(name, "ggml_backend_comm_allreduce_tensor") == 0) {
         return (void *)ggml_backend_cuda_comm_allreduce_tensor;
     }
+#ifdef GGML_USE_NCCL
+    if (strcmp(name, "ggml_backend_comm_init_rank") == 0) {
+        return (void *)ggml_backend_cuda_comm_init_rank;
+    }
+#endif // GGML_USE_NCCL
     if (strcmp(name, "ggml_backend_register_host_buffer") == 0) {
         return (void *)ggml_backend_cuda_register_host_buffer;
     }

ggml/src/ggml-cuda/quantize.cu

@@ -83,21 +83,21 @@ static __global__ void quantize_q8_1(
     ggml_cuda_pdl_sync();
     const float xi = i0 < ne00 ? x[i03*s03 + i02*s02 + i01*s01 + i00] : 0.0f;
     float amax = fabsf(xi);
-    float sum = xi;
 
     amax = warp_reduce_max<QK8_1>(amax);
-    sum  = warp_reduce_sum<QK8_1>(sum);
 
     const float  d = amax / 127.0f;
     const int8_t q = amax == 0.0f ? 0 : roundf(xi / d);
 
+    const int sum = warp_reduce_sum<QK8_1>((int) q);
+
     y[ib].qs[iqs] = q;
 
     if (iqs > 0) {
         return;
     }
 
-    y[ib].ds = make_half2(d, sum);
+    y[ib].ds = make_half2(d, d * sum);
 }
 
 __device__ __forceinline__ uint8_t compute_e8m0_scale(float amax) {
@@ -500,17 +500,6 @@ static __global__ void quantize_mmq_q8_1(
         amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, offset, WARP_SIZE));
     }
 
-    float sum;
-    if (ds_layout != MMQ_Q8_1_DS_LAYOUT_D4) {
-        sum = xi.x + xi.y + xi.z + xi.w;
-
-        // Calculate sums across vals_per_sum/4 threads.
-#pragma unroll
-        for (int offset = vals_per_sum/8; offset > 0; offset >>= 1) {
-            sum += __shfl_xor_sync(0xFFFFFFFF, sum, offset, WARP_SIZE);
-        }
-    }
-
     const float d_inv = 127.0f / amax;
     char4 q;
     q.x = roundf(xi.x*d_inv);
@@ -519,6 +508,19 @@ static __global__ void quantize_mmq_q8_1(
     q.w = roundf(xi.w*d_inv);
     const float d = 1.0f / d_inv;
 
+    float sum = 0.0f;
+    if (ds_layout != MMQ_Q8_1_DS_LAYOUT_D4) {
+        int sum_q = int(q.x) + int(q.y) + int(q.z) + int(q.w);
+
+        // Calculate sums across vals_per_sum/4 threads.
+#pragma unroll
+        for (int offset = vals_per_sum/8; offset > 0; offset >>= 1) {
+            sum_q += __shfl_xor_sync(0xFFFFFFFF, sum_q, offset, WARP_SIZE);
+        }
+
+        sum = d * float(sum_q);
+    }
+
     // write the block once (normal) or to each of the token's compact rows (scatter)
     const int nwrite = scatter ? n_expert_used : 1;
 #pragma unroll

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

@@ -1273,7 +1273,7 @@ static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggm
         case GGML_OP_FLASH_ATTN_EXT:
             if (op->type == GGML_TYPE_F32 && 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) && op->src[2] &&
-                (op->src[2]->type == GGML_TYPE_F32 || op->src[2]->type == GGML_TYPE_F16) && op->src[4] == nullptr &&
+                (op->src[2]->type == GGML_TYPE_F32 || op->src[2]->type == GGML_TYPE_F16) && op->src[4] == nullptr && op->src[5] == nullptr &&
                 ggml_is_contiguous_rows(op) && ggml_is_contiguous_rows(op->src[0])) {
                 float max_bias      = 0.0f;
                 float logit_softcap = 0.0f;

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

@@ -348,6 +348,11 @@ static bool htp_kparams_solve_flash_attn(
         return false;
     }
 
+    // K/V row lists (ggml_flash_attn_ext_rows) are not supported
+    if (op->src[5] != nullptr) {
+        return false;
+    }
+
     memset(kparams, 0, sizeof(*kparams));
 
     const struct ggml_tensor * q     = op->src[0];
@@ -2655,6 +2660,129 @@ static bool htp_kparams_solve_concat(
     return false;
 }
 
+// Permuted copy with strided dim 0 (e.g. channel-first <-> channel-last)
+static bool htp_kparams_solve_cpy_transpose(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * dst,
+    uint32_t type_size,
+    struct htp_copy_kernel_params * kparams) {
+
+    // View contiguous dst with src shape to support CPY with reshape
+    const bool dst_cont = ggml_is_contiguous(dst);
+    if (!dst_cont && !ggml_are_same_shape(src0, dst)) {
+        return false;
+    }
+
+    // Drop size-1 dims and merge dims adjacent in both tensors
+    size_t ne[GGML_MAX_DIMS];
+    size_t nbs[GGML_MAX_DIMS];
+    size_t nbd[GGML_MAX_DIMS];
+    size_t dst_stride = type_size;
+    int    n          = 0;
+    for (int d = 0; d < GGML_MAX_DIMS; d++) {
+        const size_t sd = dst_cont ? dst_stride : dst->nb[d];
+        dst_stride *= (size_t) src0->ne[d];
+        if (src0->ne[d] == 1) {
+            continue;
+        }
+        if (n > 0 && src0->nb[d] == nbs[n - 1] * ne[n - 1] && sd == nbd[n - 1] * ne[n - 1]) {
+            ne[n - 1] *= (size_t) src0->ne[d];
+            continue;
+        }
+        ne[n]  = (size_t) src0->ne[d];
+        nbs[n] = src0->nb[d];
+        nbd[n] = sd;
+        n++;
+    }
+
+    // Fold leading dense dim into element block
+    size_t es = type_size;
+    for (int i = 0; i < n; i++) {
+        if (nbs[i] == type_size && nbd[i] == type_size) {
+            es *= ne[i];
+            for (int j = i; j + 1 < n; j++) {
+                ne[j]  = ne[j + 1];
+                nbs[j] = nbs[j + 1];
+                nbd[j] = nbd[j + 1];
+            }
+            n--;
+            break;
+        }
+    }
+    if (es > 128 || (es & (es - 1)) != 0) {
+        return false;
+    }
+
+    // Transpose axes: dim a dense in dst, dim k dense in src
+    int a = -1;
+    int k = -1;
+    for (int i = 0; i < n; i++) {
+        if (a < 0 && nbd[i] == es) {
+            a = i;
+        }
+        if (k < 0 && nbs[i] == es) {
+            k = i;
+        }
+    }
+    if (a < 0 || k < 0 || a == k) {
+        return false;
+    }
+
+    struct htp_copy_transpose_params tp = {};
+    tp.elem_size = (uint8_t) es;
+    tp.ne[0]     = (uint32_t) ne[a];
+    tp.ne[1]     = (uint32_t) ne[k];
+    tp.ne[2]     = 1;
+    tp.src_nb[0] = (uint32_t) nbs[a];
+    tp.dst_nb[0] = (uint32_t) nbd[k];
+    tp.packed    = nbd[k] == ne[a] * es && ne[a] * es < 128 && 128 % (ne[a] * es) == 0;
+
+    size_t ne3 = 1;
+    int    m   = 0;
+    for (int i = 0; i < n; i++) {
+        if (i == a || i == k) {
+            continue;
+        }
+        if (m == 0) {
+            tp.ne[2] = (uint32_t) ne[i];
+        } else {
+            ne3 = ne[i];
+        }
+        tp.src_nb[1 + m] = (uint32_t) nbs[i];
+        tp.dst_nb[1 + m] = (uint32_t) nbd[i];
+        m++;
+    }
+
+    const uint32_t n_threads = htp_session_n_threads(sess);
+    const size_t   vtcm_size = htp_session_vtcm_size(sess);
+    uint64_t       vtcm      = 0;
+    for (uint32_t blk_bytes = HTP_COPY_TRANSPOSE_BLK_BYTES; ; blk_bytes /= 2) {
+        if (!htp_copy_transpose_layout_build(&tp, blk_bytes)) {
+            return false;
+        }
+        vtcm = (uint64_t) n_threads * 2 * (tp.in_buf_size + tp.out_buf_size);
+        if (vtcm <= vtcm_size) {
+            break;
+        }
+    }
+
+    const uint64_t total_tiles = (uint64_t) tp.n0b * tp.nkb * tp.ne[2] * ne3;
+    if (total_tiles > UINT32_MAX) {
+        return false;
+    }
+    tp.total_tiles = (uint32_t) total_tiles;
+    tp.div_nkb     = init_fastdiv_values(tp.nkb);
+    tp.div_n0b     = init_fastdiv_values(tp.n0b);
+    tp.div_ne2     = init_fastdiv_values(tp.ne[2]);
+
+    kparams->kernel_type = HTP_COPY_KERNEL_TRANSPOSE;
+    kparams->n_threads   = (uint8_t) n_threads;
+    kparams->vtcm_size   = (uint32_t) vtcm;
+    kparams->u.transpose = tp;
+    return true;
+}
+
 static bool htp_kparams_solve_cpy(
     const struct ggml_hexagon_session * sess,
     const struct ggml_tensor * op,
@@ -2752,6 +2880,10 @@ static bool htp_kparams_solve_cpy(
             }
         }
 
+        if (htp_kparams_solve_cpy_transpose(sess, src0, dst, src_type_size, kparams)) {
+            return true;
+        }
+
         if (sameshape) {
             kparams->kernel_type = HTP_COPY_KERNEL_SAMESHAPE_SAMETYPE;
             kparams->total_rows  = (uint32_t) (src0->ne[1] * src0->ne[2] * src0->ne[3]);

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

@@ -35,6 +35,10 @@ struct htp_copy_context {
     struct htp_copy_convert_vtcm_layout layout;
     uint8_t * vtcm_src0;
     uint8_t * vtcm_dst;
+
+    uint32_t tile_start;
+    uint32_t ntiles;
+    uint32_t tiles_per_thread;
 };
 
 #define cpy_preamble                              \
@@ -162,6 +166,171 @@ static inline void cpy_cursor_advance(struct cpy_row_cursor * c, uint32_t n, uin
     }
 }
 
+struct cpy_tile_cursor {
+    uint32_t ibk;
+    uint32_t ib0;
+    uint32_t i2;
+    uint32_t i3;
+};
+
+static inline void cpy_tile_cursor_advance(struct cpy_tile_cursor * c, const struct htp_copy_transpose_params * tp) {
+    if (++c->ibk == tp->nkb) {
+        c->ibk = 0;
+        if (++c->ib0 == tp->n0b) {
+            c->ib0 = 0;
+            if (++c->i2 == tp->ne[2]) {
+                c->i2 = 0;
+                c->i3++;
+            }
+        }
+    }
+}
+
+// Gather offsets for element e = l >> log2_m at (e % r, e / r)
+static inline HVX_Vector cpy_transpose_offsets(const struct htp_copy_transpose_params * tp) {
+    const uint32_t r_mask = (1u << tp->log2_r) - 1;
+    if (tp->elem_size != 2) {
+        const uint32_t m_mask = (1u << tp->log2_m) - 1;
+        int32_t offsets[32] __attribute__((aligned(128)));
+        for (uint32_t l = 0; l < 32; l++) {
+            const uint32_t e = l >> tp->log2_m;
+            offsets[l] = (e & r_mask) * tp->in_stride + (e >> tp->log2_r) * tp->elem_size + (l & m_mask) * 4;
+        }
+        return *(HVX_Vector *) offsets;
+    }
+    int16_t offsets[64] __attribute__((aligned(128)));
+    for (uint32_t l = 0; l < 64; l++) {
+        offsets[l] = (l & r_mask) * tp->in_stride + (l >> tp->log2_r) * 2;
+    }
+    return *(HVX_Vector *) offsets;
+}
+
+static inline void cpy_transpose_push_in(dma_queue * q, const struct htp_copy_transpose_params * tp, dma_addr_t src,
+                                         const struct cpy_tile_cursor * c, uint8_t * buf) {
+    const uint32_t a0  = c->ib0 * tp->blk0;
+    const uint32_t ak  = c->ibk * tp->blkk;
+    const uint32_t off = a0 * tp->src_nb[0] + ak * tp->elem_size + c->i2 * tp->src_nb[1] + c->i3 * tp->src_nb[2];
+    dma_queue_push(q, dma_make_data(buf, src + off), tp->in_stride, tp->src_nb[0],
+                   MIN(tp->blkk, tp->ne[1] - ak) * tp->elem_size, MIN(tp->blk0, tp->ne[0] - a0));
+}
+
+// Pipelined tile worker: DMA in -> vgather transpose -> DMA out
+static void cpy_thread_transpose(unsigned int nth, unsigned int ith, void * data) {
+    struct htp_copy_context * ct = (struct htp_copy_context *) data;
+    struct htp_ops_context * octx = ct->octx;
+    const struct htp_copy_transpose_params * tp = &ct->kparams->u.transpose;
+
+    const uint32_t t0 = ct->tile_start + ith * ct->tiles_per_thread;
+    const uint32_t t1 = MIN(t0 + ct->tiles_per_thread, ct->tile_start + ct->ntiles);
+    if (t0 >= t1) {
+        return;
+    }
+
+    dma_queue * dma_q = octx->ctx->dma[ith];
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
+    const dma_addr_t src = octx->src[0]->data;
+    const dma_addr_t dst = octx->dst->data;
+
+    const uint32_t es         = tp->elem_size;
+    const uint32_t ne0        = tp->ne[0];
+    const uint32_t nek        = tp->ne[1];
+    const uint32_t in_stride  = tp->in_stride;
+    const uint32_t out_stride = tp->out_stride;
+
+    const HVX_Vector offsets   = cpy_transpose_offsets(tp);
+    const uint32_t packed_mu   = (ne0 - 1) * in_stride + 127;
+    const uint32_t packed_step = 128 >> tp->log2_r;
+    const uint32_t epv         = 1u << tp->log2_r;
+
+    uint8_t * in_buf[2];
+    uint8_t * out_buf[2];
+    in_buf[0]  = ct->vtcm_src0 + ith * 2 * tp->in_buf_size;
+    in_buf[1]  = in_buf[0] + tp->in_buf_size;
+    out_buf[0] = ct->vtcm_dst + ith * 2 * tp->out_buf_size;
+    out_buf[1] = out_buf[0] + tp->out_buf_size;
+
+    struct cpy_tile_cursor cur;
+    {
+        const uint32_t q0 = fastdiv(t0, &tp->div_nkb);
+        const uint32_t q1 = fastdiv(q0, &tp->div_n0b);
+        cur.ibk = t0 - q0 * tp->nkb;
+        cur.ib0 = q0 - q1 * tp->n0b;
+        cur.i3  = fastdiv(q1, &tp->div_ne2);
+        cur.i2  = q1 - cur.i3 * tp->ne[2];
+    }
+
+    struct cpy_tile_cursor nxt = cur;
+
+    cpy_transpose_push_in(dma_q, tp, src, &cur, in_buf[0]);
+
+    for (uint32_t t = t0, i = 0; t < t1; t++, i++) {
+        while (dma_queue_depth(dma_q) > (i > 0 ? 1u : 0u)) {
+            dma_queue_pop(dma_q);
+        }
+
+        uint8_t * in  = in_buf[i & 1];
+        uint8_t * out = out_buf[i & 1];
+
+        const uint32_t a0 = cur.ib0 * tp->blk0;
+        const uint32_t ak = cur.ibk * tp->blkk;
+        const uint32_t n0 = MIN(tp->blk0, ne0 - a0);
+        const uint32_t nk = MIN(tp->blkk, nek - ak);
+
+        if (t + 1 < t1) {
+            cpy_tile_cursor_advance(&nxt, tp);
+            cpy_transpose_push_in(dma_q, tp, src, &nxt, in_buf[(i + 1) & 1]);
+        }
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) t);
+        if (tp->packed) {
+            const uint32_t nvec = (nk * ne0 * es + 127) >> 7;
+            if (es != 2) {
+                #pragma unroll(4)
+                for (uint32_t v = 0; v < nvec; v++) {
+                    Q6_vgather_ARMVw((HVX_Vector *) (out + v * 128), (size_t) (in + v * packed_step), packed_mu, offsets);
+                }
+            } else {
+                #pragma unroll(4)
+                for (uint32_t v = 0; v < nvec; v++) {
+                    Q6_vgather_ARMVh((HVX_Vector *) (out + v * 128), (size_t) (in + v * packed_step), packed_mu, offsets);
+                }
+            }
+        } else {
+            for (uint32_t rb = 0; rb < n0; rb += epv) {
+                const uint32_t  mu = (MIN(epv, n0 - rb) - 1) * in_stride + es - 1;
+                uint8_t *       d  = out + rb * es;
+                const uint8_t * s  = in + rb * in_stride;
+                if (es != 2) {
+                    #pragma unroll(4)
+                    for (uint32_t ik = 0; ik < nk; ik++) {
+                        Q6_vgather_ARMVw((HVX_Vector *) (d + ik * out_stride), (size_t) (s + ik * es), mu, offsets);
+                    }
+                } else {
+                    #pragma unroll(4)
+                    for (uint32_t ik = 0; ik < nk; ik++) {
+                        Q6_vgather_ARMVh((HVX_Vector *) (d + ik * out_stride), (size_t) (s + ik * 2), mu, offsets);
+                    }
+                }
+            }
+        }
+        hvx_gather_sync(out);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) t);
+
+        const uint32_t doff = a0 * es + ak * tp->dst_nb[0] + cur.i2 * tp->dst_nb[1] + cur.i3 * tp->dst_nb[2];
+        if (tp->packed) {
+            const uint32_t n = nk * ne0 * es;
+            dma_queue_push(dma_q, dma_make_data(dst + doff, out), n, n, n, 1);
+        } else {
+            dma_queue_push(dma_q, dma_make_data(dst + doff, out), tp->dst_nb[0], out_stride, n0 * es, nk);
+        }
+
+        cur = nxt;
+    }
+
+    dma_queue_flush(dma_q);
+}
+
 #define DEFINE_CPY_CONVERT_SAMESHAPE(NAME, CONV_FUNC)                                        \
 static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, void * data) { \
     struct htp_copy_context * ct = (struct htp_copy_context *) data;                         \
@@ -439,6 +608,48 @@ static int cpy_reshape(struct htp_ops_context * octx, const struct htp_copy_kern
     return HTP_STATUS_OK;
 }
 
+static int cpy_transpose(struct htp_ops_context * octx, const struct htp_copy_kernel_params * kparams) {
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    const struct htp_copy_transpose_params * tp = &kparams->u.transpose;
+
+    uint32_t tile_start = 0;
+    uint32_t ntiles     = tp->total_tiles;
+
+    if (octx->ctx->mdev.count > 1) {
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
+            tp->total_tiles, 1, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        tile_start = range.start;
+        ntiles     = range.count;
+    }
+
+    if (ntiles == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    if (kparams->vtcm_size > octx->ctx->vtcm_size) {
+        return HTP_STATUS_VTCM_TOO_SMALL;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
+
+    struct htp_copy_context ct;
+    ct.octx             = octx;
+    ct.kparams          = kparams;
+    ct.tile_start       = tile_start;
+    ct.ntiles           = ntiles;
+    ct.tiles_per_thread = fastdiv(ntiles + n_threads - 1, &octx->n_threads_div);
+
+    uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
+    ct.vtcm_src0 = vtcm_base;
+    ct.vtcm_dst  = vtcm_base + (size_t) n_threads * 2 * tp->in_buf_size;
+
+    work_queue_run(octx->ctx->work_queue, cpy_thread_transpose, &ct, n_threads);
+    return HTP_STATUS_OK;
+}
+
 int op_cpy(struct htp_ops_context * octx) {
     const struct htp_copy_kernel_params * kparams = (const struct htp_copy_kernel_params *) octx->kernel_params;
     int status = HTP_STATUS_OK;
@@ -459,6 +670,9 @@ int op_cpy(struct htp_ops_context * octx) {
         case HTP_COPY_KERNEL_RESHAPE:
             status = cpy_reshape(octx, kparams);
             break;
+        case HTP_COPY_KERNEL_TRANSPOSE:
+            status = cpy_transpose(octx, kparams);
+            break;
         case HTP_COPY_KERNEL_GATHER_ROWS:
             status = htp_gather_rows(octx, octx->src[0], NULL, &kparams->u.gather, kparams->n_threads);
             break;

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

@@ -15,6 +15,7 @@ enum htp_copy_kernel_type {
     HTP_COPY_KERNEL_RESHAPE            = 4,
     HTP_COPY_KERNEL_SCALAR             = 5,
     HTP_COPY_KERNEL_GATHER_ROWS        = 6,
+    HTP_COPY_KERNEL_TRANSPOSE          = 7,
 };
 
 struct htp_copy_convert_params {
@@ -32,6 +33,29 @@ struct htp_copy_reshape_params {
     struct fastdiv_values div_ne02_ne01_ne00;
 };
 
+// Batched 2D transpose of elem_size elements (dim 0 dense in dst, dim 1 dense in src)
+struct htp_copy_transpose_params {
+    uint32_t ne[3];          // dims 0..2
+    uint32_t src_nb[3];      // src strides for dims 0, 2, 3
+    uint32_t dst_nb[3];      // dst strides for dims 1, 2, 3
+    struct fastdiv_values div_nkb;
+    struct fastdiv_values div_n0b;
+    struct fastdiv_values div_ne2;
+    uint32_t blk0;           // dim-0 elements per tile
+    uint32_t blkk;           // dim-1 elements per tile
+    uint32_t n0b;            // tiles along dim 0
+    uint32_t nkb;            // tiles along dim 1
+    uint32_t in_stride;      // VTCM row pitch of staged src tile
+    uint32_t out_stride;     // VTCM row pitch of transposed tile
+    uint32_t in_buf_size;
+    uint32_t out_buf_size;
+    uint32_t total_tiles;
+    uint8_t  elem_size;
+    uint8_t  packed;         // gather straight to dst order
+    uint8_t  log2_m;         // log2(gather lanes per element)
+    uint8_t  log2_r;         // log2(dim-0 elements per vector)
+};
+
 struct htp_copy_kernel_params {
     uint8_t  kernel_type;
     uint8_t  src0_type_size;
@@ -46,6 +70,7 @@ struct htp_copy_kernel_params {
         struct htp_copy_convert_params convert;
         struct htp_copy_reshape_params reshape;
         struct htp_gather_rows_params  gather;
+        struct htp_copy_transpose_params transpose;
     } u;
 };
 
@@ -116,6 +141,76 @@ static inline void htp_copy_convert_solve_layout(
     htp_copy_convert_vtcm_layout_build(layout, ne00, src_type_size, dst_type_size, blk_rows, n_threads);
 }
 
+#define HTP_COPY_TRANSPOSE_BLK_BYTES (64 * 1024)
+
+static inline bool htp_copy_transpose_layout_build(struct htp_copy_transpose_params * tp, uint32_t blk_bytes) {
+    const uint32_t ne0  = tp->ne[0];
+    const uint32_t nek  = tp->ne[1];
+    const uint32_t es   = tp->elem_size;
+    const uint32_t epv  = 128 / es;
+
+    uint32_t blk0;
+    uint32_t blkk;
+    if (tp->packed) {
+        // Packed: vector holds epv / ne0 whole dst rows
+        blk0 = ne0;
+        blkk = blk_bytes / (ne0 * es);
+        if (es == 2 && blkk > 16384 / ne0 - 64) {
+            blkk = 16384 / ne0 - 64;
+        }
+        if (blkk > nek) {
+            blkk = nek;
+        }
+        if (blkk == 0) {
+            return false;
+        }
+    } else if (nek <= epv) {
+        // Stage short contiguous dim and expand dim 0
+        blkk = nek;
+        blk0 = blk_bytes / (nek * es);
+        blk0 -= blk0 % epv;
+        if (blk0 > hex_round_up(ne0, epv)) {
+            blk0 = hex_round_up(ne0, epv);
+        }
+        if (blk0 < epv) {
+            return false;
+        }
+    } else {
+        blk0 = hex_round_up(ne0, epv);
+        if (blk0 > 256) {
+            blk0 = 256;
+        }
+        blkk = blk_bytes / (blk0 * es);
+        blkk -= blkk % epv;
+        // Cap f16 blkk to 256 for 16-bit gather offsets
+        const uint32_t max_blkk = (es == 2 ? 512 : 4096) / es;
+        if (blkk > max_blkk) {
+            blkk = max_blkk;
+        }
+        if (blkk > hex_round_up(nek, epv)) {
+            blkk = hex_round_up(nek, epv);
+        }
+        if (blkk < epv) {
+            return false;
+        }
+    }
+
+    tp->blk0         = blk0;
+    tp->blkk         = blkk;
+    tp->n0b          = (ne0 + blk0 - 1) / blk0;
+    tp->nkb          = (nek + blkk - 1) / blkk;
+    // Pad row pitch by 1 element to avoid VTCM bank conflicts
+    const uint32_t pad = (es == 2) ? 2 : 4;
+    tp->in_stride    = blkk * es + pad;
+    tp->out_stride   = blk0 * es;
+    // Extra vector for packed gather overhang
+    tp->in_buf_size  = hex_round_up(blk0 * tp->in_stride + (tp->packed ? 128 : 0), 128);
+    tp->out_buf_size = hex_round_up(blkk * tp->out_stride, 128);
+    tp->log2_m       = es == 2 ? 0 : __builtin_ctz(es) - 2;
+    tp->log2_r       = __builtin_ctz(tp->packed ? ne0 : epv);
+    return true;
+}
+
 #if defined(__cplusplus)
 static_assert(sizeof(struct htp_copy_kernel_params) <= 128, "htp_copy_kernel_params is too large for kernel_params blob");
 #else

ggml/src/ggml-hexagon/htp/dma-queue.c

@@ -141,43 +141,9 @@ bool dma_queue_push_fallback_2d(dma_queue * q, dma_data ddata, size_t dst_stride
     dma_ring_flush(r1);
     r0->tail = r1->tail;
 
-    return dma_ring_push_single_2d(r0, ddata, 0, 0, 0, /*nrows=*/ 0);
-}
-
-bool dma_queue_push_fallback_contig(dma_queue * q, dma_data ddata, size_t total) {
-    dma_ring * r0 = q->ring0;
-    dma_ring * r1 = q->ring1;
-
-    if (((r0->push_idx + 1) & r0->idx_mask) == r0->pop_idx) {
-        return false;
-    }
-
-    r1->tail = r0->tail;
-
-    size_t rem_bytes   = total;
-    dma_addr_t cur_dst = ddata.dst;
-    dma_addr_t cur_src = ddata.src;
-
-    while (rem_bytes > 0) {
-        const uint32_t cur_bytes = MIN(rem_bytes, DMA_SAFE_CHUNK_SIZE);
-        dma_data cur_data = dma_make_data(cur_dst, cur_src);
-        if (!dma_ring_push_single_contig(r1, cur_data, cur_bytes)) {
-            dma_ring_flush(r1);
-            dma_ring_push_single_contig(r1, cur_data, cur_bytes);
-        }
-        cur_dst   += cur_bytes;
-        cur_src   += cur_bytes;
-        rem_bytes -= cur_bytes;
-    }
-
-    dma_ring_flush(r1);
-    r0->tail = r1->tail;
-
     return dma_ring_push_single_1d(r0, ddata, /*size=*/ 0);
 }
 
-#if __HVX_ARCH__ < 75
-
 bool dma_queue_push_fallback_1d(dma_queue * q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
     dma_ring * r0 = q->ring0;
     dma_ring * r1 = q->ring1;
@@ -188,19 +154,28 @@ bool dma_queue_push_fallback_1d(dma_queue * q, dma_data ddata, size_t dst_stride
 
     r1->tail = r0->tail;
 
-    size_t rem_rows    = nrows;
-    dma_addr_t cur_dst = ddata.dst;
-    dma_addr_t cur_src = ddata.src;
+    dma_addr_t row_dst = ddata.dst;
+    dma_addr_t row_src = ddata.src;
 
-    while (rem_rows > 0) {
-        dma_data cur_data = dma_make_data(cur_dst, cur_src);
-        if (!dma_ring_push_single_1d(r1, cur_data, row_size)) {
-            dma_ring_flush(r1);
-            dma_ring_push_single_1d(r1, cur_data, row_size);
+    for (size_t i = 0; i < nrows; i++) {
+        size_t     rem_bytes = row_size;
+        dma_addr_t cur_dst   = row_dst;
+        dma_addr_t cur_src   = row_src;
+
+        while (rem_bytes > 0) {
+            const uint32_t cur_bytes = MIN(rem_bytes, DMA_SAFE_CHUNK_SIZE);
+            dma_data cur_data = dma_make_data(cur_dst, cur_src);
+            if (!dma_ring_push_single_contig(r1, cur_data, cur_bytes)) {
+                dma_ring_flush(r1);
+                dma_ring_push_single_contig(r1, cur_data, cur_bytes);
+            }
+            cur_dst   += cur_bytes;
+            cur_src   += cur_bytes;
+            rem_bytes -= cur_bytes;
         }
-        cur_dst  += dst_stride;
-        cur_src  += src_stride;
-        rem_rows -= 1;
+
+        row_dst += dst_stride;
+        row_src += src_stride;
     }
 
     dma_ring_flush(r1);
@@ -208,5 +183,3 @@ bool dma_queue_push_fallback_1d(dma_queue * q, dma_data ddata, size_t dst_stride
 
     return dma_ring_push_single_1d(r0, ddata, /*size=*/ 0);
 }
-
-#endif

ggml/src/ggml-hexagon/htp/dma-queue.h

@@ -148,10 +148,7 @@ dma_queue_t dma_queue_alias_init(void * ptr, dma_queue_t main_q);
 void        dma_queue_alias_free(dma_queue_t q);
 
 bool        dma_queue_push_fallback_2d(dma_queue * q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows);
-bool        dma_queue_push_fallback_contig(dma_queue * q, dma_data ddata, size_t total);
-#if __HVX_ARCH__ < 75
 bool        dma_queue_push_fallback_1d(dma_queue * q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows);
-#endif
 
 // TODO: technically we don't need these and could use Q6_dmstart/wait/etc instead
 // but those do not seem to always compiler properly.
@@ -400,7 +397,7 @@ static inline bool dma_queue_push(dma_queue * q, dma_data ddata, size_t dst_stri
         if (__builtin_expect(total <= DMA_MAX_SIZE_24B, 1)) {
             return dma_ring_push_single_contig(q->ring0, ddata, total);
         }
-        return dma_queue_push_fallback_contig(q, ddata, total);
+        return dma_queue_push_fallback_1d(q, ddata, 0, 0, total, 1);
     }
 
     // 2. Hot path: Standard strided 2D (10-20% of calls)
@@ -410,14 +407,10 @@ static inline bool dma_queue_push(dma_queue * q, dma_data ddata, size_t dst_stri
     }
 
     // 3. Cold path: Descriptor chunking fallbacks (< 0.1%)
-#if __HVX_ARCH__ < 75
-    if (row_size <= DMA_MAX_SIZE_16B && (src_stride | dst_stride) <= DMA_MAX_STRIDE_16B) {
+    if ((row_size | src_stride | dst_stride) <= DMA_MAX_2D_ROW_SIZE) {
         return dma_queue_push_fallback_2d(q, ddata, dst_stride, src_stride, row_size, nrows);
     }
     return dma_queue_push_fallback_1d(q, ddata, dst_stride, src_stride, row_size, nrows);
-#else
-    return dma_queue_push_fallback_2d(q, ddata, dst_stride, src_stride, row_size, nrows);
-#endif
 }
 
 static inline void dma_sync_read(dma_queue * dma_q, void * dst, dma_addr_t src, size_t bytes) {

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

@@ -1987,6 +1987,23 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext(
     return res;
 }
 
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_rows(
+        ggml_metal_library_t lib,
+        const ggml_tensor * op) {
+    assert(op->op == GGML_OP_FLASH_ATTN_EXT);
+
+    const char * base = "kernel_flash_attn_ext_rows";
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, base);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, base, base, nullptr);
+    }
+
+    GGML_UNUSED(op);
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec_idx(
         ggml_metal_library_t lib,
         const ggml_tensor * op) {
@@ -2016,6 +2033,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
         bool    has_scap,
         bool    has_kvpad,
         bool    has_sparse,
+        bool    has_rows,
         int32_t nqpsg,
         int32_t ne,
         int32_t nsg,
@@ -2045,7 +2063,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
             dv,
             qne_suffix);
 
-    snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_sparse=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d",
+    snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_sparse=%d_rows=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d",
             base,
             has_mask,
             has_sinks,
@@ -2053,6 +2071,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
             has_scap,
             has_kvpad,
             has_sparse,
+            has_rows,
             ns10,
             ns20,
             nsg, nwg);
@@ -2067,6 +2086,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
         ggml_metal_cv_set_bool(cv, has_scap,  FC_FLASH_ATTN_EXT_VEC + 3);
         ggml_metal_cv_set_bool(cv, has_kvpad,  FC_FLASH_ATTN_EXT_VEC + 4);
         ggml_metal_cv_set_bool(cv, has_sparse, FC_FLASH_ATTN_EXT_VEC + 5);
+        ggml_metal_cv_set_bool(cv, has_rows,   FC_FLASH_ATTN_EXT_VEC + 6);
 
         ggml_metal_cv_set_int32(cv, ns10, FC_FLASH_ATTN_EXT_VEC + 20);
         ggml_metal_cv_set_int32(cv, ns20, FC_FLASH_ATTN_EXT_VEC + 21);

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

@@ -219,6 +219,10 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att
         int32_t ns10,
         int32_t ns20);
 
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_rows(
+        ggml_metal_library_t lib,
+        const struct ggml_tensor * op);
+
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec_idx(
         ggml_metal_library_t lib,
         const struct ggml_tensor * op);
@@ -232,6 +236,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att
         bool    has_scap,
         bool    has_kvpad,
         bool    has_sparse,
+        bool    has_rows,
         int32_t nqpsg,
         int32_t ne,
         int32_t nsg,

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

@@ -1727,6 +1727,39 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
         case GGML_OP_ROLL:
             return ggml_is_contiguous(op->src[0]);
         case GGML_OP_FLASH_ATTN_EXT:
+            // K/V row lists (ggml_flash_attn_ext_rows): one list of cache rows per q slice
+            if (op->src[5] != NULL) {
+                const struct ggml_tensor * kv_rows = op->src[5];
+
+                if (kv_rows->type != GGML_TYPE_I32 ||
+                    !ggml_is_contiguous(kv_rows) ||
+                    kv_rows->ne[2] != 1 || kv_rows->ne[3] != 1 ||
+                    kv_rows->ne[1] != op->src[0]->ne[3] ||
+                    !op->src[3] || op->src[1]->ne[3] != 1 || op->src[2]->ne[3] != 1) {
+                    return false;
+                }
+
+                // the non-vec kernels attend over a contiguous F16 copy of the referenced rows:
+                // - F16 K/V is gathered directly
+                // - quantized K/V is first dequantized with kernel_flash_attn_ext_kv_*_f16
+                // - other types keep attending to the whole cache (dense mask)
+                switch (op->src[1]->type) {
+                    case GGML_TYPE_F16:
+                    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:
+                        break;
+                    default:
+                        return false;
+                }
+
+                // the row gather copies F16 values in groups of 4
+                if (op->src[2]->ne[0] % 4 != 0) {
+                    return false;
+                }
+            }
             // for new head sizes, add checks here
             if (op->src[0]->ne[0] != 32 &&
                 op->src[0]->ne[0] != 40 &&
@@ -2102,19 +2135,11 @@ static void ggml_metal_buffer_rset_free(ggml_metal_buffer_t buf) {
 static void * ggml_metal_host_malloc(size_t n) {
     void * data = NULL;
 
-#if TARGET_OS_OSX
     kern_return_t err = vm_allocate((vm_map_t) mach_task_self(), (void *) &data, n, VM_FLAGS_ANYWHERE);
     if (err != KERN_SUCCESS) {
         GGML_LOG_ERROR("%s: error: vm_allocate failed\n", __func__);
         return NULL;
     }
-#else
-    const int result = posix_memalign((void **) &data, sysconf(_SC_PAGESIZE), n);
-    if (result != 0) {
-        GGML_LOG_ERROR("%s: error: posix_memalign failed\n", __func__);
-        return NULL;
-    }
-#endif
 
     return data;
 }
@@ -2296,11 +2321,7 @@ void ggml_metal_buffer_free(ggml_metal_buffer_t buf) {
     }
 
     if (buf->is_shared && buf->owned) {
-#if TARGET_OS_OSX
         vm_deallocate((vm_map_t)mach_task_self(), (vm_address_t)buf->all_data, buf->all_size);
-#else
-        free(buf->all_data);
-#endif
     }
 
     free(buf);

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

@@ -385,6 +385,19 @@ typedef struct {
     int32_t  nblocks;
 } ggml_metal_kargs_flash_attn_ext_kv_f16;
 
+// gather the K/V cache rows referenced by the kv_rows lists into contiguous F16 buffers
+typedef struct {
+    int32_t  ne11;      // number of rows per list (n_kv_rows)
+    int32_t  ne_12_2;   // number of K/V heads
+    int32_t  ne10;      // DK
+    int32_t  ne20;      // DV, set to 0 to skip gathering V (V is a view of K)
+    uint64_t v_off;     // byte offset of the gathered V region in dst
+    uint64_t nb11;      // K source strides (ne13 = 1: the cache is shared by all streams)
+    uint64_t nb12;
+    uint64_t nb21;      // V source strides (ne23 = 1)
+    uint64_t nb22;
+} ggml_metal_kargs_flash_attn_ext_rows;
+
 typedef struct {
     int32_t  ne11;
     int32_t  ne_12_2; // assume K and V are same shape

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

@@ -3261,7 +3261,10 @@ static bool ggml_metal_op_flash_attn_ext_use_tensor(const ggml_tensor * op, bool
         return false;
     }
 
-    return op->src[1]->ne[1] % OP_FLASH_ATTN_EXT_TENSOR_NCPSG == 0;
+    // with kv_rows, the attention runs over the compacted row lists (padded to 256), not over the cache
+    const int64_t ne11_attn = op->src[5] ? op->src[5]->ne[0] : op->src[1]->ne[1];
+
+    return ne11_attn % OP_FLASH_ATTN_EXT_TENSOR_NCPSG == 0;
 }
 
 // returns the n_kv_max hint if the sparse path is available for this op, or 0 otherwise
@@ -3270,6 +3273,11 @@ static bool ggml_metal_op_flash_attn_ext_use_tensor(const ggml_tensor * op, bool
 static int ggml_metal_op_flash_attn_ext_n_kv_max_sparse(const ggml_tensor * op) {
     assert(op->op == GGML_OP_FLASH_ATTN_EXT);
 
+    // K/V row lists are handled by the dense kernels, not by the sparse path
+    if (op->src[5] != nullptr) {
+        return 0;
+    }
+
     int32_t n_kv_max = 0;
     memcpy(&n_kv_max, ((const int32_t *) op->op_params) + 4, sizeof(n_kv_max));
 
@@ -3513,6 +3521,32 @@ size_t ggml_metal_op_flash_attn_ext_extra_idx(const ggml_tensor * op) {
     return GGML_PAD(sizeof(int32_t)*(size_t) n_kv_max_padded*ne31*ne32*ne33, 16);
 }
 
+// size of the contiguous F16 buffers holding the K/V rows referenced by the kv_rows lists
+// the attention over these buffers is a standard dense op (no per-tile gather in the main kernel)
+size_t ggml_metal_op_flash_attn_ext_extra_rows(const ggml_tensor * op) {
+    assert(op->op == GGML_OP_FLASH_ATTN_EXT);
+
+    const ggml_tensor * kv_rows = op->src[5];
+
+    if (!kv_rows) {
+        return 0;
+    }
+
+    GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne);
+    GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne);
+
+    const int64_t n_rows   = kv_rows->ne[0];
+    const int64_t n_slices = kv_rows->ne[1];
+
+    // note: when V is a view of K, the gathered V is read from the gathered K, so only K is gathered
+    const bool v_is_view_of_k = ggml_metal_op_flash_attn_ext_v_is_view_of_k(op);
+
+    const size_t k_size = GGML_PAD(sizeof(ggml_fp16_t)*(size_t) ne10*n_rows*ne12*n_slices, 16);
+    const size_t v_size = v_is_view_of_k ? 0 : GGML_PAD(sizeof(ggml_fp16_t)*(size_t) ne20*n_rows*ne22*n_slices, 16);
+
+    return k_size + v_size;
+}
+
 int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
     ggml_tensor * op = ctx->node(idx);
 
@@ -3545,6 +3579,22 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
     GGML_ASSERT(!op->src[3] || op->src[3]->ne[1] >= op->src[0]->ne[1] &&
             "the Flash-Attention Metal kernel requires the mask to be at least n_queries big");
 
+    // K/V row lists (ggml_flash_attn_ext_rows): src[5] holds one list of cache rows per q slice
+    // the attention loops then run over the mask columns and read K/V rows through the lists
+    const ggml_tensor * kv_rows = op->src[5];
+    const bool has_rows = kv_rows != nullptr;
+
+    GGML_ASSERT(!has_rows || (kv_rows->type == GGML_TYPE_I32 && ggml_is_contiguous(kv_rows)));
+    GGML_ASSERT(!has_rows || (kv_rows->ne[1] == ne03 && kv_rows->ne[2] == 1 && kv_rows->ne[3] == 1));
+    GGML_ASSERT(!has_rows || op->src[3]);
+
+    // the loop bound over the mask columns / K/V row lists
+    const int64_t ne11_attn = has_rows ? kv_rows->ne[0] : ne11;
+
+    // the row lists are padded by the KV cache to a multiple of 256, so the pad path never triggers
+    GGML_ASSERT(!has_rows || ne11_attn % OP_FLASH_ATTN_EXT_NCPSG == 0);
+    GGML_ASSERT(!has_rows || ne11_attn % OP_FLASH_ATTN_EXT_VEC_NCPSG == 0);
+
     float scale;
     float max_bias;
     float logit_softcap;
@@ -3559,6 +3609,8 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
 
     const bool has_mask  = op->src[3] != NULL;
     const bool has_sinks = op->src[4] != NULL;
+
+    GGML_ASSERT(!has_rows || has_mask);
     const bool has_bias  = max_bias != 0.0f;
     const bool has_scap  = logit_softcap != 0.0f;
 
@@ -3595,12 +3647,25 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
     const bool use_sparse = n_kv_max_sparse > 0;
     const int n_kv_max_padded = use_sparse ? GGML_PAD(n_kv_max_sparse, OP_FLASH_ATTN_EXT_VEC_NCPSG) : 0;
 
+    const bool use_vec = ggml_metal_op_flash_attn_ext_use_vec(op);
+
+    // with kv_rows, the non-vec path gathers the referenced rows into a contiguous F16 buffer and
+    // attends over it with the F16 kernel - quantized KV is first dequantized with the F16 pass
+    const bool rows_gather = has_rows && !use_vec;
+
     // the vec kernels dequantize the KV inline; no need for the F16 dequant pass in the sparse path
-    const bool use_kv_f16 = !use_sparse && ggml_metal_op_flash_attn_ext_use_kv_f16(op);
+    const bool dequant_kv = !use_sparse && (rows_gather ?
+        ggml_is_quantized(op->src[1]->type) : ggml_metal_op_flash_attn_ext_use_kv_f16(op));
+
+    const bool use_kv_f16 = !use_sparse && (rows_gather ? true : dequant_kv);
 
     ggml_metal_buffer_id bid_idx = bid_kv_f16;
     bid_idx.offs += ggml_metal_op_flash_attn_ext_extra_kv_f16(op);
 
+    // contiguous K/V buffers holding the rows referenced by the kv_rows lists
+    ggml_metal_buffer_id bid_rows = bid_idx;
+    bid_rows.offs += ggml_metal_op_flash_attn_ext_extra_idx(op);
+
     ggml_metal_buffer_id bid_k = bid_src1;
     ggml_metal_buffer_id bid_v = bid_src2;
 
@@ -3613,7 +3678,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
     uint64_t nb22_attn = nb22;
     uint64_t nb23_attn = nb23;
 
-    if (use_kv_f16) {
+    if (dequant_kv) {
         assert(ggml_metal_op_flash_attn_ext_extra_kv_f16(op) != 0);
 
         const bool v_is_view_of_k = ggml_metal_op_flash_attn_ext_v_is_view_of_k(op);
@@ -3701,6 +3766,70 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
         }
     }
 
+    // with kv_rows, gather the cache rows referenced by the lists into contiguous F16 buffers
+    // the attention then runs as a standard dense op over the gathered rows (tensor or half8x8 kernel)
+    if (rows_gather) {
+        assert(ggml_metal_op_flash_attn_ext_extra_rows(op) != 0);
+
+        const bool v_is_view_of_k = ggml_metal_op_flash_attn_ext_v_is_view_of_k(op);
+
+        GGML_ASSERT(ne20 % 4 == 0);
+
+        const size_t k_size = GGML_PAD(sizeof(ggml_fp16_t)*(size_t) ne10*ne11_attn*ne12*ne03, 16);
+
+        ggml_metal_kargs_flash_attn_ext_rows args0 = {
+            /*.ne11    =*/ (int32_t) ne11_attn,
+            /*.ne_12_2 =*/ ne12,
+            /*.ne10    =*/ ne10,
+            /*.ne20    =*/ v_is_view_of_k ? 0 : ne20, // the gathered V is a view of the gathered K
+            /*.v_off   =*/ k_size,
+            /*.nb11    =*/ nb11_attn,
+            /*.nb12    =*/ nb12_attn,
+            /*.nb21    =*/ nb21_attn,
+            /*.nb22    =*/ nb22_attn,
+        };
+
+        auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_rows(lib, op);
+        const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline0), 64);
+
+        ggml_metal_encoder_set_pipeline(enc, pipeline0);
+        ggml_metal_encoder_set_bytes   (enc, &args0, sizeof(args0), 0);
+        ggml_metal_encoder_set_buffer  (enc, bid_k, 1);
+        ggml_metal_encoder_set_buffer  (enc, bid_v, 2);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(kv_rows), 3);
+        ggml_metal_encoder_set_buffer  (enc, bid_rows, 4);
+
+        ggml_metal_encoder_dispatch_threadgroups(enc, ne11_attn, ne03, 1, nth, 1, 1);
+
+        ggml_metal_op_concurrency_reset(ctx);
+
+        // contiguous F16 layout of the gathered rows: [DK, n_rows, n_head_kv, n_streams]
+        bid_k = bid_rows;
+
+        nb10_attn = sizeof(ggml_fp16_t);
+        nb11_attn = nb10_attn*ne10;
+        nb12_attn = nb11_attn*ne11_attn;
+        nb13_attn = nb12_attn*ne12;
+
+        // if V is a view of K, the gathered V is read from the gathered K with K's strides
+        if (v_is_view_of_k) {
+            bid_v = bid_k;
+
+            nb20_attn = nb10_attn;
+            nb21_attn = nb11_attn;
+            nb22_attn = nb12_attn;
+            nb23_attn = nb13_attn;
+        } else {
+            bid_v = bid_rows;
+            bid_v.offs += k_size;
+
+            nb20_attn = sizeof(ggml_fp16_t);
+            nb21_attn = nb20_attn*ne20;
+            nb22_attn = nb21_attn*ne11_attn;
+            nb23_attn = nb22_attn*ne22;
+        }
+    }
+
     if (!use_sparse && ggml_metal_op_flash_attn_ext_use_tensor(op, props_dev->has_tensor)) {
         // tensor API kernel
         const int nqptg = ne00 >= 512 ? OP_FLASH_ATTN_EXT_TENSOR_NQPSG_LARGE : OP_FLASH_ATTN_EXT_TENSOR_NQPSG; // queries per threadgroup
@@ -3746,9 +3875,9 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
             /*.nb01          =*/ nb01,
             /*.nb02          =*/ nb02,
             /*.nb03          =*/ nb03,
-            /*.ne11          =*/ ne11,
+            /*.ne11          =*/ (int32_t) ne11_attn,
             /*.ne_12_2       =*/ ne12,
-            /*.ne_12_3       =*/ ne13,
+            /*.ne_12_3       =*/ has_rows ? ne03 : ne13, // the gathered K/V holds one copy per stream
             /*.ns10          =*/ ns10,
             /*.nb11          =*/ nb11_attn,
             /*.nb12          =*/ nb12_attn,
@@ -3806,7 +3935,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
 
         bool need_sync = false;
 
-        const bool has_kvpad = ne11 % ncpsg != 0;
+        const bool has_kvpad = ne11_attn % ncpsg != 0;
 
         if (has_kvpad) {
             assert(ggml_metal_op_flash_attn_ext_extra_pad(op) != 0);
@@ -3906,9 +4035,9 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
             /*.nb01          =*/ nb01,
             /*.nb02          =*/ nb02,
             /*.nb03          =*/ nb03,
-            /*.ne11          =*/ ne11,
+            /*.ne11          =*/ (int32_t) ne11_attn,
             /*.ne_12_2       =*/ ne12,
-            /*.ne_12_3       =*/ ne13,
+            /*.ne_12_3       =*/ has_rows ? ne03 : ne13, // the gathered K/V holds one copy per stream
             /*.ns10          =*/ ns10,
             /*.nb11          =*/ nb11_attn,
             /*.nb12          =*/ nb12_attn,
@@ -3974,7 +4103,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
 
         bool need_sync = false;
 
-        const bool has_kvpad = !use_sparse && ne11 % ncpsg != 0;
+        const bool has_kvpad = !use_sparse && ne11_attn % ncpsg != 0;
 
         if (use_sparse) {
             assert(ggml_metal_op_flash_attn_ext_extra_idx(op) != 0);
@@ -4057,13 +4186,13 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
         //   queries:      Q*NSG*PAD2(ne00, 128)
         //   attn + mask:  NSG*4*Q*C
         //   results:      2*NSG*Q*PAD2(ne20, 128)
-        //   sparse idx:   NSG*C ints (only when use_sparse)
+        //   sparse idx:   NSG*C ints (when use_sparse, or when gathering K/V rows)
         const int64_t dk_pad = GGML_PAD(ne00, 128);
         const int64_t dv_pad = GGML_PAD(ne20, 128);
 
         auto fa_vec_smem = [&](int64_t nsg, int32_t nqptg) -> size_t {
             const size_t smem_half = (size_t) (dk_pad + 4*ncpsg + 2*dv_pad)*nqptg*nsg;
-            return GGML_PAD(smem_half*sizeof(ggml_fp16_t) + (use_sparse ? (size_t) nsg*ncpsg*sizeof(int) : 0), 16);
+            return GGML_PAD(smem_half*sizeof(ggml_fp16_t) + ((use_sparse || has_rows) ? (size_t) nsg*ncpsg*sizeof(int) : 0), 16);
         };
 
         int64_t nsg = 1;
@@ -4094,7 +4223,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
         } else {
             nwg = 32;
             nsg = 1;
-            while (2*nwg*nsg*ncpsg < ne11 && nsg < 4) {
+            while (2*nwg*nsg*ncpsg < ne11_attn && nsg < 4) {
                 nsg *= 2;
             }
         }
@@ -4115,7 +4244,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
             /*.nb01          =*/ nb01,
             /*.nb02          =*/ nb02,
             /*.nb03          =*/ nb03,
-            /*.ne11          =*/ use_sparse ? n_kv_max_padded : ne11,
+            /*.ne11          =*/ use_sparse ? n_kv_max_padded : (int32_t) ne11_attn,
             /*.ne_12_2       =*/ ne12,
             /*.ne_12_3       =*/ ne13,
             /*.ns10          =*/ ns10,
@@ -4144,7 +4273,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
             /*.n_kv_max_padded =*/ n_kv_max_padded,
         };
 
-        auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, use_sparse, nqptg, cfg.NE, nsg, nwg, use_kv_f16, ns10, ns20);
+        auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, use_sparse, has_rows, nqptg, cfg.NE, nsg, nwg, use_kv_f16, ns10, ns20);
 
         GGML_ASSERT(nsg*32 <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
 
@@ -4155,7 +4284,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
         ggml_metal_encoder_set_buffer  (enc, bid_v,    3);
         ggml_metal_encoder_set_buffer  (enc, bid_src3, 4);
         ggml_metal_encoder_set_buffer  (enc, bid_src4, 5);
-        ggml_metal_encoder_set_buffer  (enc, use_sparse ? bid_idx : bid_src0, 8);
+        ggml_metal_encoder_set_buffer  (enc, use_sparse ? bid_idx : (has_rows ? ggml_metal_get_buffer_id(kv_rows) : bid_src0), 8);
 
         const size_t smem = fa_vec_smem(nsg, nqptg);
 

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

@@ -45,6 +45,7 @@ size_t ggml_metal_op_flash_attn_ext_extra_blk(const struct ggml_tensor * op);
 size_t ggml_metal_op_flash_attn_ext_extra_tmp(const struct ggml_tensor * op);
 size_t ggml_metal_op_flash_attn_ext_extra_kv_f16(const struct ggml_tensor * op);
 size_t ggml_metal_op_flash_attn_ext_extra_idx(const struct ggml_tensor * op);
+size_t ggml_metal_op_flash_attn_ext_extra_rows(const struct ggml_tensor * op);
 
 size_t ggml_metal_op_argmax_extra(const struct ggml_tensor * op);
 

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

@@ -300,6 +300,7 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_
                 res += ggml_metal_op_flash_attn_ext_extra_tmp(tensor);
                 res += ggml_metal_op_flash_attn_ext_extra_kv_f16(tensor);
                 res += ggml_metal_op_flash_attn_ext_extra_idx(tensor);
+                res += ggml_metal_op_flash_attn_ext_extra_rows(tensor);
             } break;
         case GGML_OP_CUMSUM:
         case GGML_OP_ARGSORT:

ggml/src/ggml-metal/kernels/fa_aux.metal

@@ -48,6 +48,80 @@ template [[host_name("kernel_flash_attn_ext_kv_q5_0_f16")]] kernel kernel_flash_
 template [[host_name("kernel_flash_attn_ext_kv_q5_1_f16")]] kernel kernel_flash_attn_ext_kv_f16_t kernel_flash_attn_ext_kv_f16<block_q5_1, 32, dequantize_q5_1>;
 template [[host_name("kernel_flash_attn_ext_kv_q8_0_f16")]] kernel kernel_flash_attn_ext_kv_f16_t kernel_flash_attn_ext_kv_f16<block_q8_0, 32, dequantize_q8_0>;
 
+// gather the K/V cache rows referenced by the kv_rows lists into contiguous F16 buffers
+// - one threadgroup per (row, stream); the attention is then run as a standard dense op over the
+//   gathered rows, which avoids the per-tile K/V gather in the main kernels
+// - negative rows are padding: filled with zeros, their mask entries are -INF so they contribute nothing
+kernel void kernel_flash_attn_ext_rows(
+        constant ggml_metal_kargs_flash_attn_ext_rows & args,
+        device const char * k_src,
+        device const char * v_src,
+        device const int  * rows,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort  tiitg[[thread_index_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    const int64_t i1 = tgpig[0]; // row index
+    const int64_t i3 = tgpig[1]; // stream
+
+    const int64_t row = rows[i1 + i3*args.ne11];
+
+    // gathered layout: [DK, n_rows, n_head_kv, n_streams] and [DV, n_rows, n_head_kv, n_streams] F16
+    device half4 * k_dst = (device half4 *) (dst + i3*args.ne_12_2*args.ne11*args.ne10*sizeof(half));
+    device half4 * v_dst = (device half4 *) (dst + args.v_off + i3*args.ne_12_2*args.ne11*args.ne20*sizeof(half));
+
+    // note: the K/V cache is shared by all streams (ne13 = ne23 = 1), the row list selects the rows
+    // each thread copies up to 4 half4 values (32 bytes) to keep the memory transactions large
+
+    // K: all heads of this row
+    {
+        const int64_t per_head = args.ne10/4;
+
+        for (int64_t i = 4*tiitg; i < args.ne_12_2*per_head; i += 4*ntg.x) {
+            const int64_t i2 = i/per_head;
+            const int64_t i0 = i%per_head;
+            const int      n  = (int) min((int64_t) 4, per_head - i0);
+
+            device half4 * dst4 = k_dst + (i2*args.ne11 + i1)*per_head + i0;
+
+            if (row >= 0) {
+                device const half4 * src4 = (device const half4 *) (k_src + row*args.nb11 + i2*args.nb12) + i0;
+                for (int j = 0; j < n; ++j) {
+                    dst4[j] = src4[j];
+                }
+            } else {
+                for (int j = 0; j < n; ++j) {
+                    dst4[j] = half4(0.0h);
+                }
+            }
+        }
+    }
+
+    // V: all heads of this row
+    {
+        const int64_t per_head = args.ne20/4;
+
+        for (int64_t i = 4*tiitg; i < args.ne_12_2*per_head; i += 4*ntg.x) {
+            const int64_t i2 = i/per_head;
+            const int64_t i0 = i%per_head;
+            const int      n  = (int) min((int64_t) 4, per_head - i0);
+
+            device half4 * dst4 = v_dst + (i2*args.ne11 + i1)*per_head + i0;
+
+            if (row >= 0) {
+                device const half4 * src4 = (device const half4 *) (v_src + row*args.nb21 + i2*args.nb22) + i0;
+                for (int j = 0; j < n; ++j) {
+                    dst4[j] = src4[j];
+                }
+            } else {
+                for (int j = 0; j < n; ++j) {
+                    dst4[j] = half4(0.0h);
+                }
+            }
+        }
+    }
+}
+
 constant bool FC_flash_attn_ext_pad_has_mask [[function_constant(FC_FLASH_ATTN_EXT_PAD + 0)]];
 
 constant int32_t FC_flash_attn_ext_pad_ncpsg [[function_constant(FC_FLASH_ATTN_EXT_PAD + 25)]];

ggml/src/ggml-metal/kernels/fa_vec_common.metal

@@ -13,6 +13,7 @@ constant int32_t FC_flash_attn_ext_vec_ns20 [[function_constant(FC_FLASH_ATTN_EX
 constant int32_t FC_flash_attn_ext_vec_nsg  [[function_constant(FC_FLASH_ATTN_EXT_VEC + 22)]];
 constant int32_t FC_flash_attn_ext_vec_nwg  [[function_constant(FC_FLASH_ATTN_EXT_VEC + 23)]];
 constant bool    FC_flash_attn_ext_vec_has_sparse [[function_constant(FC_FLASH_ATTN_EXT_VEC + 5)]];
+constant bool    FC_flash_attn_ext_vec_has_rows   [[function_constant(FC_FLASH_ATTN_EXT_VEC + 6)]];
 template<
     typename q4_t,  // query types in shared memory
     typename k4_t,  // key types in shared memory
@@ -90,7 +91,8 @@ kernel void kernel_flash_attn_ext_vec(
     threadgroup o4_t  * so4 = (threadgroup o4_t  *) (shmem_f16 + SMEM_Q + SMEM_S + 2*sgitg*Q*PV); // scratch buffer for the results
 
     // sparse indices for the current block
-    threadgroup int * spidx = FC_flash_attn_ext_vec_has_sparse
+    // note: with kv_rows, the buffer holds the row list of the current slice instead
+    threadgroup int * spidx = (FC_flash_attn_ext_vec_has_sparse || FC_flash_attn_ext_vec_has_rows)
         ? (threadgroup int *) (shmem_f16 + SMEM) + sgitg*C
         : nullptr;
 
@@ -165,6 +167,12 @@ kernel void kernel_flash_attn_ext_vec(
                 ((int64_t)(iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*args.ne31*args.n_kv_max_padded + (iq1%args.ne31)*args.n_kv_max_padded;
         }
 
+        // K/V row list of the current slice (ggml_flash_attn_ext_rows)
+        // with kv_rows, the loop below runs over the mask columns (args.ne11) and the K/V cache rows
+        // are read indirectly through the list. negative entries are padding, their mask is -INF
+        device const int * rows_list = FC_flash_attn_ext_vec_has_rows ?
+            (device const int *) idx + (int64_t)(iq3%args.ne33)*args.ne11 : nullptr;
+
         float slope = 1.0f;
 
         // ALiBi
@@ -222,12 +230,12 @@ kernel void kernel_flash_attn_ext_vec(
                 ic = 0;
             }
 
-            // load the sparse KV indices for the current block into shared memory
-            if (FC_flash_attn_ext_vec_has_sparse) {
+            // load the KV row indices for the current block into shared memory
+            if (FC_flash_attn_ext_vec_has_sparse || FC_flash_attn_ext_vec_has_rows) {
                 FOR_UNROLL (short ii = 0; ii < C/NW; ++ii) {
                     const short i = ii*NW + tiisg;
 
-                    spidx[i] = pidx[ic + i];
+                    spidx[i] = FC_flash_attn_ext_vec_has_rows ? rows_list[ic + i] : pidx[ic + i];
                 }
                 simdgroup_barrier(mem_flags::mem_threadgroup);
             }
@@ -242,6 +250,15 @@ kernel void kernel_flash_attn_ext_vec(
                             sm[qq*C + tiisg] = -MAXHALF;
                         }
                     }
+                } else if (FC_flash_attn_ext_vec_has_rows) {
+                    // the mask is compact: column ic maps to cache row spidx[tiisg]
+                    FOR_UNROLL (short qq = 0; qq < Q; ++qq) {
+                        if ((iq1*Q + qq) < args.ne01) {
+                            sm[qq*C + tiisg] = pm[qq][ic + tiisg];
+                        } else {
+                            sm[qq*C + tiisg] = -MAXHALF;
+                        }
+                    }
                 } else {
                     FOR_UNROLL (short qq = 0; qq < Q; ++qq) {
                         if ((iq1*Q + qq) < args.ne01) {
@@ -276,7 +293,7 @@ kernel void kernel_flash_attn_ext_vec(
             {
                 device      const k4_t * pk4 = nullptr;
 
-                if (!FC_flash_attn_ext_vec_has_sparse) {
+                if (!FC_flash_attn_ext_vec_has_sparse && !FC_flash_attn_ext_vec_has_rows) {
                     pk4 = (device const k4_t *) (k + ic*args.nb11);
 
                     pk4 += ty*NS10/4 + tx;
@@ -291,7 +308,7 @@ kernel void kernel_flash_attn_ext_vec(
 
                 // each simdgroup processes Q queries and NE (NW/NL) cache elements
                 FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) {
-                    if (FC_flash_attn_ext_vec_has_sparse) {
+                    if (FC_flash_attn_ext_vec_has_sparse || FC_flash_attn_ext_vec_has_rows) {
                         // the KV rows are gathered from the index list; -1 entries are padding
                         const int i11 = spidx[NE*cc + ty];
                         if (i11 >= 0) {
@@ -434,7 +451,7 @@ kernel void kernel_flash_attn_ext_vec(
                     }
                 }
 
-                if (FC_flash_attn_ext_vec_has_sparse) {
+                if (FC_flash_attn_ext_vec_has_sparse || FC_flash_attn_ext_vec_has_rows) {
                     FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) {
                         // the KV rows are gathered from the index list; -1 entries are padding
                         const int i11 = spidx[NE*cc + ty];

ggml/src/ggml-metal/kernels/quantize.metal

@@ -26,11 +26,11 @@ kernel void kernel_cpy_t_t(
     const int32_t i1 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0)/args.ne0;
     const int32_t i0 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0 - i1*args.ne0);
 
-    device T1 * dst_data = (device T1 *) (dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
+    device char * dst_data = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0;
 
     for (int32_t i00 = iw0*ntg[0] + tpitg.x; i00 < args.ne00;) {
         device const T0 * src = (device T0 *)(src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + i00*args.nb00);
-        dst_data[i00] = (T1) src[0];
+        *(device T1 *)(dst_data + i00*args.nb0) = (T1) src[0];
         break;
     }
 }

ggml/src/ggml-opencl/CMakeLists.txt

@@ -67,6 +67,7 @@ set(GGML_OPENCL_KERNELS
     argmax
     argsort
     top_k
+    top_k_mask
     lightning_indexer
     tri
     fill
@@ -124,6 +125,15 @@ set(GGML_OPENCL_KERNELS
     gemm_noshuffle_iq1_s_f32
     gemm_noshuffle_iq1_s_q8_1_dp4a
     gemv_noshuffle_iq1_s_f32
+    gemm_noshuffle_iq3_s_f32
+    gemm_noshuffle_iq3_s_q8_1_dp4a
+    gemv_noshuffle_iq3_s_f32
+    gemm_noshuffle_iq3_xxs_f32
+    gemm_noshuffle_iq3_xxs_q8_1_dp4a
+    gemv_noshuffle_iq3_xxs_f32
+    gemm_noshuffle_iq4_xs_f32
+    gemm_noshuffle_iq4_xs_q8_1_dp4a
+    gemv_noshuffle_iq4_xs_f32
     gemm_noshuffle_q3_k_q8_1_dp4a
     gemv_noshuffle_q3_k_f32
     mul_mv_q2_k_f32
@@ -150,6 +160,18 @@ set(GGML_OPENCL_KERNELS
     mul_mv_iq1_s_f32
     mul_mv_iq1_s_f32_flat
     mul_mm_iq1_s_f32_l4_lm
+    iq3_xxs_grid
+    mul_mv_iq3_xxs_f32
+    mul_mv_iq3_xxs_f32_flat
+    mul_mm_iq3_xxs_f32_l4_lm
+    iq3_s_grid
+    mul_mv_iq3_s_f32
+    mul_mv_iq3_s_f32_flat
+    mul_mm_iq3_s_f32_l4_lm
+    iq4_xs_grid
+    mul_mv_iq4_xs_f32
+    mul_mv_iq4_xs_f32_flat
+    mul_mm_iq4_xs_f32_l4_lm
     mul_mv_q3_k_f32_flat
     mul_mm_q3_k_f32_l4_lm
     mul_mv_q2_k_f32_flat

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

@@ -3178,3 +3178,203 @@ kernel void kernel_restore_block_iq1_m(
         dst[i].scales[2*j + 1] = (uchar)(sc >> 8);
     }
 }
+
+struct block_iq3_s {
+    half d;
+    uchar qs[QK_K/4];
+    uchar qh[QK_K/32];
+    uchar signs[QK_K/8];
+    uchar scales[QK_K/64];
+};
+
+kernel void kernel_convert_block_iq3_s(
+    global struct block_iq3_s * src0,
+    global uchar * dst_qs,
+    global uchar * dst_sg,
+    global uchar * dst_qh,
+    global uchar * dst_sc,
+    global half * dst_d,
+    ulong n_blocks
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blocks) { return; }
+    dst_d[i] = src0[i].d;
+    for (int j = 0; j < 64; ++j) {
+        dst_qs[64*i + j] = src0[i].qs[j];
+    }
+    for (int j = 0; j < 32; ++j) {
+        dst_sg[32*i + j] = src0[i].signs[j];
+    }
+    for (int j = 0; j < 8; ++j) {
+        dst_qh[8*i + j] = src0[i].qh[j];
+    }
+    for (int j = 0; j < 4; ++j) {
+        dst_sc[4*i + j] = src0[i].scales[j];
+    }
+}
+
+kernel void kernel_restore_block_iq3_s(
+    global uchar * src_qs,
+    global uchar * src_sg,
+    global uchar * src_qh,
+    global uchar * src_sc,
+    global half * src_d,
+    global struct block_iq3_s * dst,
+    ulong n_blocks
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blocks) { return; }
+    dst[i].d = src_d[i];
+    for (int j = 0; j < 64; ++j) {
+        dst[i].qs[j] = src_qs[64*i + j];
+    }
+    for (int j = 0; j < 32; ++j) {
+        dst[i].signs[j] = src_sg[32*i + j];
+    }
+    for (int j = 0; j < 8; ++j) {
+        dst[i].qh[j] = src_qh[8*i + j];
+    }
+    for (int j = 0; j < 4; ++j) {
+        dst[i].scales[j] = src_sc[4*i + j];
+    }
+}
+
+struct block_iq3_xxs {
+    half d;
+    uchar qs[3*QK_K/8];
+};
+
+kernel void kernel_convert_block_iq3_xxs(
+    global struct block_iq3_xxs * src0,
+    global uchar * dst_qs,
+    global uint * dst_sas,
+    global half * dst_d,
+    ulong n_blocks
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blocks) { return; }
+    dst_d[i] = src0[i].d;
+    for (int j = 0; j < 64; ++j) {
+        dst_qs[64*i + j] = src0[i].qs[j];
+    }
+    for (int j = 0; j < 8; ++j) {
+        global ushort * q = (global ushort *)(src0[i].qs + 64 + 4*j);
+        dst_sas[8*i + j] = (uint)q[0] | ((uint)q[1] << 16);
+    }
+}
+
+kernel void kernel_restore_block_iq3_xxs(
+    global uchar * src_qs,
+    global uint * src_sas,
+    global half * src_d,
+    global struct block_iq3_xxs * dst,
+    ulong n_blocks
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blocks) { return; }
+    dst[i].d = src_d[i];
+    for (int j = 0; j < 64; ++j) {
+        dst[i].qs[j] = src_qs[64*i + j];
+    }
+    for (int j = 0; j < 8; ++j) {
+        global ushort * q = (global ushort *)(dst[i].qs + 64 + 4*j);
+        uint sas = src_sas[8*i + j];
+        q[0] = (ushort)sas;
+        q[1] = (ushort)(sas >> 16);
+    }
+}
+
+struct block_iq4_xs {
+    half d;
+    ushort scales_h;
+    uchar scales_l[QK_K/64];
+    uchar qs[QK_K/2];
+};
+
+kernel void kernel_convert_block_iq4_xs(
+    global struct block_iq4_xs * src0,
+    global uchar * dst_qs,
+    global uchar * dst_sl,
+    global ushort * dst_sh,
+    global half * dst_d,
+    ulong n_blocks
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blocks) { return; }
+    dst_d[i] = src0[i].d;
+    dst_sh[i] = src0[i].scales_h;
+    for (int j = 0; j < 128; ++j) {
+        dst_qs[128*i + j] = src0[i].qs[j];
+    }
+    for (int j = 0; j < 4; ++j) {
+        dst_sl[4*i + j] = src0[i].scales_l[j];
+    }
+}
+
+kernel void kernel_restore_block_iq4_xs(
+    global uchar * src_qs,
+    global uchar * src_sl,
+    global ushort * src_sh,
+    global half * src_d,
+    global struct block_iq4_xs * dst,
+    ulong n_blocks
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blocks) { return; }
+    dst[i].d = src_d[i];
+    dst[i].scales_h = src_sh[i];
+    for (int j = 0; j < 128; ++j) {
+        dst[i].qs[j] = src_qs[128*i + j];
+    }
+    for (int j = 0; j < 4; ++j) {
+        dst[i].scales_l[j] = src_sl[4*i + j];
+    }
+}
+
+kernel void kernel_convert_block_iq4_xs_noshuffle(
+    global struct block_iq4_xs * src0,
+    global ushort * dst_qs,
+    global uint * dst_sl,
+    global ushort * dst_sh,
+    global half * dst_d,
+    ulong n_blocks
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blocks) { return; }
+    dst_d[i] = src0[i].d;
+    dst_sh[i] = src0[i].scales_h;
+    dst_sl[i] = as_uint(vload4(0, src0[i].scales_l));
+    for (int sb = 0; sb < 8; ++sb) {
+        for (int u = 0; u < 8; ++u) {
+            ushort q = 0;
+            for (int j = 0; j < 4; ++j) {
+                uchar v = src0[i].qs[16*sb + 4*(u & 3) + j];
+                q |= (ushort)(((v >> (4*(u >> 2))) & 15) << (4*j));
+            }
+            dst_qs[64*i + 8*sb + u] = q;
+        }
+    }
+}
+
+kernel void kernel_restore_block_iq4_xs_noshuffle(
+    global ushort * src_qs,
+    global uint * src_sl,
+    global ushort * src_sh,
+    global half * src_d,
+    global struct block_iq4_xs * dst,
+    ulong n_blocks
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blocks) { return; }
+    dst[i].d = src_d[i];
+    dst[i].scales_h = src_sh[i];
+    vstore4(as_uchar4(src_sl[i]), 0, dst[i].scales_l);
+    for (int sb = 0; sb < 8; ++sb) {
+        for (int j = 0; j < 16; ++j) {
+            const int shift = 4*(j & 3);
+            const ushort lo = src_qs[64*i + 8*sb + j/4];
+            const ushort hi = src_qs[64*i + 8*sb + j/4 + 4];
+            dst[i].qs[16*sb + j] = (uchar)(((lo >> shift) & 15) | (((hi >> shift) & 15) << 4));
+        }
+    }
+}

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

@@ -0,0 +1,92 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_qcom_reqd_sub_group_size
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+__attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+kernel void kernel_gemm_noshuffle_iq3_s_f32(
+    global const uchar * src0_qs,
+    global const uchar * src0_sg,
+    global const uchar * src0_qh,
+    global const uchar * src0_sc,
+    global const half  * src0_d,
+    read_only image1d_buffer_t src1,
+    global float * dst,
+    ulong offsetd,
+    int m,
+    int n,
+    int k,
+    int n_no_padding
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+    const int gy = get_global_id(0);
+    const int row = get_global_id(1) * 4;
+    const int n_4 = n / 4;
+    if (row >= m) {
+        return;
+    }
+
+    half8 c0 = 0, c1 = 0, c2 = 0, c3 = 0;
+    for (int i = 0; i < k; i += 16) {
+        const uchar4 sc = vload4(0, src0_sc + row + (i / 64) * m);
+        const half4 d = vload4(0, src0_d + row + (i / 256) * m);
+        const float4 scale_f = convert_float4(d) * (1.0f + 2.0f*convert_float4((sc >> (uchar4)(4*((i / 32) & 1))) & (uchar4)(15)));
+        const half4 scale = convert_half4(scale_f);
+
+        for (int l = 0; l < 16; l += 4) {
+            const int ki = i + l;
+            const uint4 qs = convert_uint4(vload4(0, src0_qs + row + (ki / 4) * m));
+            const uint4 qh = convert_uint4(vload4(0, src0_qh + row + (ki / 32) * m));
+            const uchar4 sg = vload4(0, src0_sg + row + (ki / 8) * m);
+            const uint4 gi = qs | (((qh >> (uint4)(((ki / 4) & 7))) & (uint4)(1)) << (uint4)(8));
+            const uint4 grid = (uint4)(iq3_s_grid[gi.s0], iq3_s_grid[gi.s1],
+                                      iq3_s_grid[gi.s2], iq3_s_grid[gi.s3]);
+            #pragma unroll
+            for (int j = 0; j < 4; ++j) {
+                half8 b;
+                b.s0123 = read_imageh(src1, gy * 2     + (ki + j) * n_4);
+                b.s4567 = read_imageh(src1, gy * 2 + 1 + (ki + j) * n_4);
+                const int4 magnitude = convert_int4((grid >> (uint4)(8*j)) & (uint4)(255));
+                const int4 signs = convert_int4((sg >> (uchar4)((ki & 7) + j)) & (uchar4)(1));
+                const half4 w = convert_half4(magnitude * ((int4)(1) - 2*signs)) * scale;
+                c0 += b * w.s0;
+                c1 += b * w.s1;
+                c2 += b * w.s2;
+                c3 += b * w.s3;
+            }
+        }
+    }
+
+    int idx = gy * 8 * m + row;
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s0, c1.s0, c2.s0, c3.s0), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s1, c1.s1, c2.s1, c3.s1), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s2, c1.s2, c2.s2, c3.s2), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s3, c1.s3, c2.s3, c3.s3), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s4, c1.s4, c2.s4, c3.s4), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s5, c1.s5, c2.s5, c3.s5), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s6, c1.s6, c2.s6, c3.s6), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s7, c1.s7, c2.s7, c3.s7), 0, dst + idx);
+    }
+}

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

@@ -0,0 +1,154 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+#ifdef cl_khr_integer_dot_product
+#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
+#endif
+
+#define QK_K 256
+
+#ifndef TILESIZE_N
+#define TILESIZE_N 32
+#endif
+
+inline uint iq3_s_pack(uint gv, uint sg, uint base) {
+    int v0 = (int)((gv >>  0) & 0xFF); if (sg & (1u << (base + 0))) { v0 = -v0; }
+    int v1 = (int)((gv >>  8) & 0xFF); if (sg & (1u << (base + 1))) { v1 = -v1; }
+    int v2 = (int)((gv >> 16) & 0xFF); if (sg & (1u << (base + 2))) { v2 = -v2; }
+    int v3 = (int)((gv >> 24) & 0xFF); if (sg & (1u << (base + 3))) { v3 = -v3; }
+    return ((uint)v0 & 0xFFu) | (((uint)v1 & 0xFFu) <<  8)
+         | (((uint)v2 & 0xFFu) << 16) | (((uint)v3 & 0xFFu) << 24);
+}
+
+#define IQ3_S_GEMM_GRID(i) iq3_s_grid[(i)]
+
+inline uint iq3_s_load(__global const uchar * qs, uint qsb, uint qhb, uint m,
+                      uint u, uint sgv, uint base) {
+    const uint gidx = (uint)qs[qsb + u * m] | (((qhb >> u) & 1u) << 8);
+    return iq3_s_pack(IQ3_S_GEMM_GRID(gidx), sgv, base);
+}
+
+// Pass staged operands by value; local dp4a operands miscompile on X2.
+inline int dot8_q8a_v(uint8 qw, uint4 a0, uint4 a1) {
+    int r = 0;
+    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a0.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a0.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a0.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a0.w, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s4, a1.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s5, a1.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s6, a1.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s7, a1.w, r);
+    return r;
+}
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_iq3_s_q8_1_dp4a(
+        __global const uchar  * src0_qs,
+        __global const uchar  * src0_sg,
+        __global const uchar  * src0_qh,
+        __global const uchar  * src0_sc,
+        __global const half   * src0_d,
+        __global const uint   * src1_qa,
+        __global const half   * src1_da,
+        __global       float  * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local half sh_d[TILESIZE_N];
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub = step >> 5;
+
+        const uint ib = sub >> 3;
+        const uint sb = sub & 7u;
+
+        const uint  scb = (uint)src0_sc[rrow + (ib*(QK_K/64) + (sb >> 1)) * (uint)m];
+        const uint  nib = (sb & 1u) ? (scb >> 4) : (scb & 0xFu);
+        const float d_w = (float)src0_d[rrow + ib * (uint)m] * (float)(1 + 2*nib);
+
+        const uint qhb  = (uint)src0_qh[rrow + (ib*(QK_K/32) + sb) * (uint)m];
+        const uint qsb  = rrow + (ib*(QK_K/4) + 8*sb) * (uint)m;
+        const uint sgb  = rrow + (ib*(QK_K/8) + 4*sb) * (uint)m;
+
+        const uint sg0 = (uint)src0_sg[sgb + 0 * (uint)m];
+        const uint sg1 = (uint)src0_sg[sgb + 1 * (uint)m];
+        const uint sg2 = (uint)src0_sg[sgb + 2 * (uint)m];
+        const uint sg3 = (uint)src0_sg[sgb + 3 * (uint)m];
+
+        uint8 qw;
+        qw.s0 = iq3_s_load(src0_qs, qsb, qhb, (uint)m, 0u, sg0, 0u);
+        qw.s1 = iq3_s_load(src0_qs, qsb, qhb, (uint)m, 1u, sg0, 4u);
+        qw.s2 = iq3_s_load(src0_qs, qsb, qhb, (uint)m, 2u, sg1, 0u);
+        qw.s3 = iq3_s_load(src0_qs, qsb, qhb, (uint)m, 3u, sg1, 4u);
+        qw.s4 = iq3_s_load(src0_qs, qsb, qhb, (uint)m, 4u, sg2, 0u);
+        qw.s5 = iq3_s_load(src0_qs, qsb, qhb, (uint)m, 5u, sg2, 4u);
+        qw.s6 = iq3_s_load(src0_qs, qsb, qhb, (uint)m, 6u, sg3, 0u);
+        qw.s7 = iq3_s_load(src0_qs, qsb, qhb, (uint)m, 7u, sg3, 4u);
+
+        for (uint idx = lid; idx < TILESIZE_N * 2; idx += 64) {
+            const uint t = idx >> 1;
+            const uint v = idx & 1;
+            const uint c = col_base + t;
+            sh_qa4[t][v] = (c < (uint)n_no_padding)
+                         ? vload4(0, src1_qa + c * k_u + (step >> 2) + (v << 2))
+                         : (uint4)(0u);
+        }
+        if (lid < TILESIZE_N) {
+            const uint c = col_base + lid;
+            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+#define TDOT4(T) dot8_q8a_v(qw, (uint4)(sh_qa4[T][0]), (uint4)(sh_qa4[T][1]))
+            rf.s0 = (float)TDOT4(b+0);  rf.s1 = (float)TDOT4(b+1);
+            rf.s2 = (float)TDOT4(b+2);  rf.s3 = (float)TDOT4(b+3);
+            acc[g] += d_w * LD4(sh_d, b) * rf;
+        }
+#undef TDOT4
+#undef LD4
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
+        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
+        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
+        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
+    }
+#undef NGROUPS
+}

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

@@ -0,0 +1,89 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_qcom_reqd_sub_group_size
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+__attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+kernel void kernel_gemm_noshuffle_iq3_xxs_f32(
+    global const uchar * src0_qs,
+    global const uint * src0_sas,
+    global const half  * src0_d,
+    read_only image1d_buffer_t src1,
+    global float * dst,
+    ulong offsetd,
+    int m,
+    int n,
+    int k,
+    int n_no_padding
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+    const int gy = get_global_id(0);
+    const int row = get_global_id(1) * 4;
+    const int n_4 = n / 4;
+    if (row >= m) {
+        return;
+    }
+
+    half8 c0 = 0, c1 = 0, c2 = 0, c3 = 0;
+    for (int i = 0; i < k; i += 16) {
+        const uint4 sas = vload4(0, src0_sas + row + (i / 32) * m);
+        const half4 d = vload4(0, src0_d + row + (i / 256) * m);
+        const float4 scale_f = convert_float4(d) * 0.5f * (0.5f + convert_float4(sas >> (uint4)(28)));
+        const half4 scale = convert_half4(scale_f);
+
+        for (int l = 0; l < 16; l += 4) {
+            const int ki = i + l;
+            const uint4 gi = convert_uint4(vload4(0, src0_qs + row + (ki / 4) * m));
+            const uint4 code = (sas >> (uint4)(7*((ki / 8) & 3))) & (uint4)(127);
+            const uint4 sg = code | ((popcount(code) & (uint4)(1)) << (uint4)(7));
+            const uint4 grid = (uint4)(iq3_xxs_grid[gi.s0], iq3_xxs_grid[gi.s1],
+                                      iq3_xxs_grid[gi.s2], iq3_xxs_grid[gi.s3]);
+            #pragma unroll
+            for (int j = 0; j < 4; ++j) {
+                half8 b;
+                b.s0123 = read_imageh(src1, gy * 2     + (ki + j) * n_4);
+                b.s4567 = read_imageh(src1, gy * 2 + 1 + (ki + j) * n_4);
+                const int4 magnitude = convert_int4((grid >> (uint4)(8*j)) & (uint4)(255));
+                const int4 signs = convert_int4((sg >> (uint4)((ki & 7) + j)) & (uint4)(1));
+                const half4 w = convert_half4(magnitude * ((int4)(1) - 2*signs)) * scale;
+                c0 += b * w.s0;
+                c1 += b * w.s1;
+                c2 += b * w.s2;
+                c3 += b * w.s3;
+            }
+        }
+    }
+
+    int idx = gy * 8 * m + row;
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s0, c1.s0, c2.s0, c3.s0), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s1, c1.s1, c2.s1, c3.s1), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s2, c1.s2, c2.s2, c3.s2), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s3, c1.s3, c2.s3, c3.s3), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s4, c1.s4, c2.s4, c3.s4), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s5, c1.s5, c2.s5, c3.s5), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s6, c1.s6, c2.s6, c3.s6), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s7, c1.s7, c2.s7, c3.s7), 0, dst + idx);
+    }
+}

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

@@ -0,0 +1,137 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+#ifdef cl_khr_integer_dot_product
+#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
+#endif
+
+#define QK_K 256
+
+#ifndef TILESIZE_N
+#define TILESIZE_N 32
+#endif
+
+inline uint iq3_xxs_pack(uint gv, uint sg, uint base) {
+    int v0 = (int)((gv >>  0) & 0xFF); if (sg & (1u << (base + 0))) { v0 = -v0; }
+    int v1 = (int)((gv >>  8) & 0xFF); if (sg & (1u << (base + 1))) { v1 = -v1; }
+    int v2 = (int)((gv >> 16) & 0xFF); if (sg & (1u << (base + 2))) { v2 = -v2; }
+    int v3 = (int)((gv >> 24) & 0xFF); if (sg & (1u << (base + 3))) { v3 = -v3; }
+    return ((uint)v0 & 0xFFu) | (((uint)v1 & 0xFFu) <<  8)
+         | (((uint)v2 & 0xFFu) << 16) | (((uint)v3 & 0xFFu) << 24);
+}
+
+// Pass staged operands by value; local dp4a operands miscompile on X2.
+inline int dot8_q8a_v(uint8 qw, uint4 a0, uint4 a1) {
+    int r = 0;
+    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a0.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a0.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a0.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a0.w, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s4, a1.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s5, a1.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s6, a1.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s7, a1.w, r);
+    return r;
+}
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_iq3_xxs_q8_1_dp4a(
+        __global const uchar  * src0_qs,
+        __global const uint   * src0_sas,
+        __global const half   * src0_d,
+        __global const uint   * src1_qa,
+        __global const half   * src1_da,
+        __global       float  * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local half sh_d[TILESIZE_N];
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub = step >> 5;
+
+        const uint ib = sub >> 3;
+        const uint aux = src0_sas[rrow + sub*(uint)m];
+        const float d_w = (float)src0_d[rrow + ib*(uint)m] * (0.5f + (float)(aux >> 28)) * 0.5f;
+        const uint sg0 = iq3_xxs_signs( aux        & 127u);
+        const uint sg1 = iq3_xxs_signs((aux >>  7) & 127u);
+        const uint sg2 = iq3_xxs_signs((aux >> 14) & 127u);
+        const uint sg3 = iq3_xxs_signs((aux >> 21) & 127u);
+        const uint qsb = rrow + sub*8u*(uint)m;
+
+        uint8 qw;
+        qw.s0 = iq3_xxs_pack(iq3_xxs_grid[src0_qs[qsb + 0u*(uint)m]], sg0, 0u);
+        qw.s1 = iq3_xxs_pack(iq3_xxs_grid[src0_qs[qsb + 1u*(uint)m]], sg0, 4u);
+        qw.s2 = iq3_xxs_pack(iq3_xxs_grid[src0_qs[qsb + 2u*(uint)m]], sg1, 0u);
+        qw.s3 = iq3_xxs_pack(iq3_xxs_grid[src0_qs[qsb + 3u*(uint)m]], sg1, 4u);
+        qw.s4 = iq3_xxs_pack(iq3_xxs_grid[src0_qs[qsb + 4u*(uint)m]], sg2, 0u);
+        qw.s5 = iq3_xxs_pack(iq3_xxs_grid[src0_qs[qsb + 5u*(uint)m]], sg2, 4u);
+        qw.s6 = iq3_xxs_pack(iq3_xxs_grid[src0_qs[qsb + 6u*(uint)m]], sg3, 0u);
+        qw.s7 = iq3_xxs_pack(iq3_xxs_grid[src0_qs[qsb + 7u*(uint)m]], sg3, 4u);
+
+        for (uint idx = lid; idx < TILESIZE_N * 2; idx += 64) {
+            const uint t = idx >> 1;
+            const uint v = idx & 1;
+            const uint c = col_base + t;
+            sh_qa4[t][v] = (c < (uint)n_no_padding)
+                         ? vload4(0, src1_qa + c * k_u + (step >> 2) + (v << 2))
+                         : (uint4)(0u);
+        }
+        if (lid < TILESIZE_N) {
+            const uint c = col_base + lid;
+            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+#define TDOT4(T) dot8_q8a_v(qw, (uint4)(sh_qa4[T][0]), (uint4)(sh_qa4[T][1]))
+            rf.s0 = (float)TDOT4(b+0);  rf.s1 = (float)TDOT4(b+1);
+            rf.s2 = (float)TDOT4(b+2);  rf.s3 = (float)TDOT4(b+3);
+            acc[g] += d_w * LD4(sh_d, b) * rf;
+        }
+#undef TDOT4
+#undef LD4
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
+        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
+        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
+        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
+    }
+#undef NGROUPS
+}

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

@@ -0,0 +1,88 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_qcom_reqd_sub_group_size
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+__attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+kernel void kernel_gemm_noshuffle_iq4_xs_f32(
+    global const ushort * src0_qs,
+    global const uint * src0_sl,
+    global const ushort * src0_sh,
+    global const half  * src0_d,
+    read_only image1d_buffer_t src1,
+    global float * dst,
+    ulong offsetd,
+    int m,
+    int n,
+    int k,
+    int n_no_padding
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+    const int gy = get_global_id(0);
+    const int row = get_global_id(1) * 4;
+    const int n_4 = n / 4;
+    if (row >= m) {
+        return;
+    }
+
+    half8 c0 = 0, c1 = 0, c2 = 0, c3 = 0;
+    for (int i = 0; i < k; i += 16) {
+        const uint4 sl = vload4(0, src0_sl + row + (i / 256) * m);
+        const uint4 sh = convert_uint4(vload4(0, src0_sh + row + (i / 256) * m));
+        const half4 d = vload4(0, src0_d + row + (i / 256) * m);
+        const int sb = (i / 32) & 7;
+        const int4 ls = convert_int4(((sl >> (uint4)(4*sb)) & (uint4)(15)) | (((sh >> (uint4)(2*sb)) & (uint4)(3)) << (uint4)(4))) - (int4)(32);
+        const half4 scale = convert_half4(convert_float4(d) * convert_float4(ls));
+
+        for (int l = 0; l < 16; l += 4) {
+            const int ki = i + l;
+            const ushort4 q = vload4(0, src0_qs + row + (ki / 4) * m);
+            #pragma unroll
+            for (int j = 0; j < 4; ++j) {
+                half8 b;
+                b.s0123 = read_imageh(src1, gy * 2     + (ki + j) * n_4);
+                b.s4567 = read_imageh(src1, gy * 2 + 1 + (ki + j) * n_4);
+                const ushort4 indices = (q >> (ushort4)(4*j)) & (ushort4)(15);
+                const half4 w = (half4)(iq4_xs_values[indices.s0], iq4_xs_values[indices.s1],
+                                       iq4_xs_values[indices.s2], iq4_xs_values[indices.s3]) * scale;
+                c0 += b * w.s0;
+                c1 += b * w.s1;
+                c2 += b * w.s2;
+                c3 += b * w.s3;
+            }
+        }
+    }
+
+    int idx = gy * 8 * m + row;
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s0, c1.s0, c2.s0, c3.s0), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s1, c1.s1, c2.s1, c3.s1), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s2, c1.s2, c2.s2, c3.s2), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s3, c1.s3, c2.s3, c3.s3), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s4, c1.s4, c2.s4, c3.s4), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s5, c1.s5, c2.s5, c3.s5), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s6, c1.s6, c2.s6, c3.s6), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s7, c1.s7, c2.s7, c3.s7), 0, dst + idx);
+    }
+}

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

@@ -0,0 +1,135 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+#ifdef cl_khr_integer_dot_product
+#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
+#endif
+
+#define QK_K 256
+
+#ifndef TILESIZE_N
+#define TILESIZE_N 32
+#endif
+
+inline uint iq4_xs_pack(ushort q) {
+    return iq4_xs_code(q & 15)
+         | (iq4_xs_code((q >> 4) & 15) << 8)
+         | (iq4_xs_code((q >> 8) & 15) << 16)
+         | (iq4_xs_code(q >> 12) << 24);
+}
+
+// Pass staged operands by value; local dp4a operands miscompile on X2.
+inline int dot8_q8a_v(uint8 qw, uint4 a0, uint4 a1) {
+    int r = 0;
+    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a0.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a0.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a0.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a0.w, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s4, a1.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s5, a1.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s6, a1.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s7, a1.w, r);
+    return r;
+}
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_iq4_xs_q8_1_dp4a(
+        __global const ushort * src0_qs,
+        __global const uint   * src0_sl,
+        __global const ushort * src0_sh,
+        __global const half   * src0_d,
+        __global const uint   * src1_qa,
+        __global const half   * src1_da,
+        __global       float  * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local half sh_d[TILESIZE_N];
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub = step >> 5;
+
+        const uint ib = sub >> 3;
+        const uint sb = sub & 7u;
+        const uint sl = src0_sl[rrow + ib*(uint)m];
+        const uint sh = src0_sh[rrow + ib*(uint)m];
+        const int ls = (int)(((sl >> (4*sb)) & 15u) | (((sh >> (2*sb)) & 3u) << 4)) - 32;
+        const float d_w = (float)src0_d[rrow + ib*(uint)m] * (float)ls;
+        const uint qsb = rrow + sub*8u*(uint)m;
+
+        uint8 qw;
+        qw.s0 = iq4_xs_pack(src0_qs[qsb + 0u*(uint)m]);
+        qw.s1 = iq4_xs_pack(src0_qs[qsb + 1u*(uint)m]);
+        qw.s2 = iq4_xs_pack(src0_qs[qsb + 2u*(uint)m]);
+        qw.s3 = iq4_xs_pack(src0_qs[qsb + 3u*(uint)m]);
+        qw.s4 = iq4_xs_pack(src0_qs[qsb + 4u*(uint)m]);
+        qw.s5 = iq4_xs_pack(src0_qs[qsb + 5u*(uint)m]);
+        qw.s6 = iq4_xs_pack(src0_qs[qsb + 6u*(uint)m]);
+        qw.s7 = iq4_xs_pack(src0_qs[qsb + 7u*(uint)m]);
+
+        for (uint idx = lid; idx < TILESIZE_N * 2; idx += 64) {
+            const uint t = idx >> 1;
+            const uint v = idx & 1;
+            const uint c = col_base + t;
+            sh_qa4[t][v] = (c < (uint)n_no_padding)
+                         ? vload4(0, src1_qa + c * k_u + (step >> 2) + (v << 2))
+                         : (uint4)(0u);
+        }
+        if (lid < TILESIZE_N) {
+            const uint c = col_base + lid;
+            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+#define TDOT4(T) dot8_q8a_v(qw, (uint4)(sh_qa4[T][0]), (uint4)(sh_qa4[T][1]))
+            rf.s0 = (float)TDOT4(b+0);  rf.s1 = (float)TDOT4(b+1);
+            rf.s2 = (float)TDOT4(b+2);  rf.s3 = (float)TDOT4(b+3);
+            acc[g] += d_w * LD4(sh_d, b) * rf;
+        }
+#undef TDOT4
+#undef LD4
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
+        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
+        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
+        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
+    }
+#undef NGROUPS
+}

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

@@ -0,0 +1,128 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#define QK_K 256
+
+#ifndef IQ3_S_MV_NSG
+#define IQ3_S_MV_NSG 4
+#endif
+
+inline float4 iq3_s_vals(uint gv, uint sg, uint base) {
+
+    const uint  s   = sg >> base;
+    const uint4 sgn = (uint4)(s << 31, s << 30, s << 29, s << 28) & 0x80000000u;
+    return as_float4(as_uint4(convert_float4(as_uchar4(gv))) ^ sgn);
+}
+
+kernel void kernel_gemv_noshuffle_iq3_s_f32(
+        __read_only image1d_buffer_t grid_img,
+        global const uchar * src0_qs,
+        global const uchar * src0_sg,
+        global const uchar * src0_qh,
+        global const uchar * src0_sc,
+        global const half  * src0_d,
+        global const float * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        int ne10,
+        int ne0
+) {
+    src1 = (global const float *)((global const char *)src1 + offset1);
+    dst  = (global float       *)((global char       *)dst  + offsetd);
+
+    const uint m   = (uint)ne01;
+    const uint K   = (uint)ne00;
+    const uint nsb = K / QK_K;
+
+    const uint lid = get_local_id(0);
+    const uint sgi = get_local_id(1);
+    const uint col = get_group_id(1);
+
+    global const float * y = src1 + (ulong)col * (uint)ne10;
+
+#define IQ3_S_GRID(i) (read_imageui(grid_img, (int)(i)).x)
+
+    const uint mh  = m >> 1;
+    const uint j   = get_group_id(0) * 64u + lid;
+    const uint row = j << 1;
+
+    float sumf  = 0.f;
+    float sumf1 = 0.f;
+
+    if (j < mh) {
+        global const ushort * qsu = (global const ushort *)src0_qs;
+        global const ushort * qhu = (global const ushort *)src0_qh;
+        global const ushort * sgu = (global const ushort *)src0_sg;
+        global const ushort * scu = (global const ushort *)src0_sc;
+
+        for (uint ib = sgi; ib < nsb; ib += IQ3_S_MV_NSG) {
+            const half2 dh = vload2(j + ib * mh, src0_d);
+            const float d0 = (float)dh.s0;
+            const float d1 = (float)dh.s1;
+
+            const uint scbase = j + ib * 4u * mh;
+            ushort sc4[4];
+            sc4[0] = scu[scbase + 0u * mh];
+            sc4[1] = scu[scbase + 1u * mh];
+            sc4[2] = scu[scbase + 2u * mh];
+            sc4[3] = scu[scbase + 3u * mh];
+
+            float acc0 = 0.f, acc1 = 0.f;
+            for (uint sb = 0; sb < 8u; ++sb) {
+                const uint scv  = (uint)sc4[sb >> 1];
+                const uint nib0 = (sb & 1u) ? ((scv >>  4) & 0xFu) : ( scv        & 0xFu);
+                const uint nib1 = (sb & 1u) ? ((scv >> 12) & 0xFu) : ((scv >>  8) & 0xFu);
+
+                const uint qhv = (uint)qhu[j + (ib * 8u + sb) * mh];
+                const uint qh0 = qhv & 0xFFu;
+                const uint qh1 = qhv >> 8;
+
+                const uint grp = ib * 64u + sb * 8u;
+                const uint qsb = j + grp * mh;
+                const uint sgb = j + (ib * 32u + sb * 4u) * mh;
+
+                float a0 = 0.f, a1 = 0.f;
+                for (uint u = 0; u < 8u; ++u) {
+                    const uint qsv = (uint)qsu[qsb + u * mh];
+                    const uint sgv = (uint)sgu[sgb + (u >> 1) * mh];
+                    const uint g0  = ( qsv       & 0xFFu) | (((qh0 >> u) & 1u) << 8);
+                    const uint g1  = ((qsv >> 8) & 0xFFu) | (((qh1 >> u) & 1u) << 8);
+                    const uint base = (u & 1u) * 4u;
+                    const float4 yv = vload4(grp + u, y);
+                    const uint gw0 = IQ3_S_GRID(g0);
+                    const uint gw1 = IQ3_S_GRID(g1);
+                    a0 += dot(yv, iq3_s_vals(gw0,  sgv       & 0xFFu, base));
+                    a1 += dot(yv, iq3_s_vals(gw1, (sgv >> 8) & 0xFFu, base));
+                }
+                acc0 += (float)(1u + 2u * nib0) * a0;
+                acc1 += (float)(1u + 2u * nib1) * a1;
+            }
+            sumf  += d0 * acc0;
+            sumf1 += d1 * acc1;
+        }
+    }
+
+#if IQ3_S_MV_NSG > 1
+    __local float2 part[IQ3_S_MV_NSG][64];
+    part[sgi][lid] = (float2)(sumf, sumf1);
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (sgi != 0) {
+        return;
+    }
+    for (uint s = 1; s < IQ3_S_MV_NSG; ++s) {
+        const float2 p = part[s][lid];
+        sumf  += p.s0;
+        sumf1 += p.s1;
+    }
+#endif
+
+    if (j < mh) {
+        vstore2((float2)(sumf, sumf1), 0, dst + (ulong)col * (uint)ne0 + row);
+    }
+}
+
+kernel void kernel_iq3_s_grid_export(global uint * out) {
+    out[get_global_id(0)] = iq3_s_grid[get_global_id(0)];
+}

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

@@ -0,0 +1,120 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#define QK_K 256
+
+#ifndef IQ3_XXS_MV_NSG
+#define IQ3_XXS_MV_NSG 4
+#endif
+
+inline float4 iq3_xxs_vals(uint gv, uint sgv, uint base) {
+    const uint  s   = sgv >> base;
+    const uint4 sgn = (uint4)(s << 31, s << 30, s << 29, s << 28) & 0x80000000u;
+    return as_float4(as_uint4(convert_float4(as_uchar4(gv))) ^ sgn);
+}
+
+kernel void kernel_gemv_noshuffle_iq3_xxs_f32(
+        __read_only image1d_buffer_t grid_img,
+        global const uchar * src0_qs,
+        global const uint  * src0_sas,
+        global const half  * src0_d,
+        global const float * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        int ne10,
+        int ne0
+) {
+    src1 = (global const float *)((global const char *)src1 + offset1);
+    dst  = (global float       *)((global char       *)dst  + offsetd);
+
+    const uint m   = (uint)ne01;
+    const uint K   = (uint)ne00;
+    const uint nsb = K / QK_K;
+
+    const uint lid = get_local_id(0);
+    const uint sgi = get_local_id(1);
+    const uint col = get_group_id(1);
+
+    global const float * y = src1 + (ulong)col * (uint)ne10;
+
+#define IQ3_XXS_GRID(i) (read_imageui(grid_img, (int)(i)).x)
+
+    const uint mh  = m >> 1;
+    const uint j   = get_group_id(0) * 64u + lid;
+    const uint row = j << 1;
+
+    float sumf  = 0.f;
+    float sumf1 = 0.f;
+
+    if (j < mh) {
+        global const ushort * qsu = (global const ushort *)src0_qs;
+
+        for (uint ib = sgi; ib < nsb; ib += IQ3_XXS_MV_NSG) {
+            const half2 dh = vload2(j + ib * mh, src0_d);
+            const float d0 = (float)dh.s0;
+            const float d1 = (float)dh.s1;
+
+            float acc0 = 0.f, acc1 = 0.f;
+            for (uint sb = 0; sb < 8u; ++sb) {
+                const uint  sub = ib * 8u + sb;
+                const uint2 aux = vload2(j + sub * mh, src0_sas);
+
+                const uint grp = sub * 8u;
+                const uint qsb = j + grp * mh;
+
+                float a0 = 0.f, a1 = 0.f;
+                for (uint u = 0; u < 4u; ++u) {
+                    const uint qsv0 = (uint)qsu[qsb + (2u*u)*mh];
+                    const uint qsv1 = (uint)qsu[qsb + (2u*u + 1u)*mh];
+                    const uint sg0 = iq3_xxs_signs((aux.s0 >> (7u * u)) & 127u);
+                    const uint sg1 = iq3_xxs_signs((aux.s1 >> (7u * u)) & 127u);
+
+                    const uint g00 =  qsv0       & 0xFFu;
+                    const uint g01 =  qsv1       & 0xFFu;
+                    const uint g10 = (qsv0 >> 8) & 0xFFu;
+                    const uint g11 = (qsv1 >> 8) & 0xFFu;
+
+                    const float4 yl = vload4(sub * 8u + u * 2u + 0u, y);
+                    const float4 yh = vload4(sub * 8u + u * 2u + 1u, y);
+
+                    a0 += dot(yl, iq3_xxs_vals(IQ3_XXS_GRID(g00), sg0, 0u))
+                        + dot(yh, iq3_xxs_vals(IQ3_XXS_GRID(g01), sg0, 4u));
+                    a1 += dot(yl, iq3_xxs_vals(IQ3_XXS_GRID(g10), sg1, 0u))
+                        + dot(yh, iq3_xxs_vals(IQ3_XXS_GRID(g11), sg1, 4u));
+                }
+                acc0 += (0.5f + (float)(aux.s0 >> 28)) * a0;
+                acc1 += (0.5f + (float)(aux.s1 >> 28)) * a1;
+            }
+            sumf  += d0 * 0.5f * acc0;
+            sumf1 += d1 * 0.5f * acc1;
+        }
+    }
+
+#if IQ3_XXS_MV_NSG > 1
+    __local float2 part[IQ3_XXS_MV_NSG][64];
+    part[sgi][lid] = (float2)(sumf, sumf1);
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (sgi != 0) {
+        return;
+    }
+    for (uint s = 1; s < IQ3_XXS_MV_NSG; ++s) {
+        const float2 p = part[s][lid];
+        sumf  += p.s0;
+        sumf1 += p.s1;
+    }
+#endif
+
+    if (j < mh) {
+        vstore2((float2)(sumf, sumf1), 0, dst + (ulong)col * (uint)ne0 + row);
+    }
+#undef IQ3_XXS_GRID
+}
+
+kernel void kernel_iq3_xxs_grid_export(global uint * out) {
+    const uint i = get_global_id(0);
+    if (i < 256u) {
+        out[i] = iq3_xxs_grid[i];
+    }
+}

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

@@ -0,0 +1,85 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#define QK_K 256
+
+#ifndef IQ4_XS_MV_NSG
+#define IQ4_XS_MV_NSG 8
+#endif
+
+inline float4 iq4_xs_vals(ushort q) {
+    return (float4)((char)iq4_xs_code(q & 15), (char)iq4_xs_code((q >> 4) & 15),
+                    (char)iq4_xs_code((q >> 8) & 15), (char)iq4_xs_code(q >> 12));
+}
+
+kernel void kernel_gemv_noshuffle_iq4_xs_f32(
+        global const ushort * src0_qs,
+        global const uint * src0_sl,
+        global const ushort * src0_sh,
+        global const half * src0_d,
+        global const float * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        int ne10,
+        int ne0
+) {
+    src1 = (global const float *)((global const char *)src1 + offset1);
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint m = (uint)ne01;
+    const uint mh = m >> 1;
+    const uint lid = get_local_id(0);
+    const uint sgi = get_local_id(1);
+    const uint col = get_group_id(1);
+    const uint j = get_group_id(0)*64u + lid;
+    const uint row = j << 1;
+    global const float * y = src1 + (ulong)col * (uint)ne10;
+
+    float sum0 = 0.0f, sum1 = 0.0f;
+    if (j < mh) {
+        global const uint * qs = (global const uint *)src0_qs;
+        global const uint * sh = (global const uint *)src0_sh;
+        for (uint ib = sgi; ib < (uint)ne00/QK_K; ib += IQ4_XS_MV_NSG) {
+            const uint base = j + ib*mh;
+            const uint2 sl = vload2(base, src0_sl);
+            const uint shp = sh[base];
+            const half2 d = vload2(base, src0_d);
+            float acc0 = 0.0f, acc1 = 0.0f;
+            for (uint sb = 0; sb < 8; ++sb) {
+                const int ls0 = (int)(((sl.s0 >> (4*sb)) & 15u) | (((shp >> (2*sb)) & 3u) << 4)) - 32;
+                const int ls1 = (int)(((sl.s1 >> (4*sb)) & 15u) | (((shp >> (16 + 2*sb)) & 3u) << 4)) - 32;
+                float a0 = 0.0f, a1 = 0.0f;
+                for (uint u = 0; u < 8; ++u) {
+                    const uint group = ib*64 + sb*8 + u;
+                    const uint q = qs[j + group*mh];
+                    const float4 yv = vload4(group, y);
+                    a0 += dot(yv, iq4_xs_vals((ushort)q));
+                    a1 += dot(yv, iq4_xs_vals((ushort)(q >> 16)));
+                }
+                acc0 += (float)ls0*a0;
+                acc1 += (float)ls1*a1;
+            }
+            sum0 += (float)d.s0*acc0;
+            sum1 += (float)d.s1*acc1;
+        }
+    }
+
+#if IQ4_XS_MV_NSG > 1
+    local float2 part[IQ4_XS_MV_NSG][64];
+    part[sgi][lid] = (float2)(sum0, sum1);
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (sgi != 0) {
+        return;
+    }
+    for (uint s = 1; s < IQ4_XS_MV_NSG; ++s) {
+        const float2 v = part[s][lid];
+        sum0 += v.s0;
+        sum1 += v.s1;
+    }
+#endif
+    if (j < mh) {
+        vstore2((float2)(sum0, sum1), 0, dst + (ulong)col*(uint)ne0 + row);
+    }
+}

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

@@ -0,0 +1,66 @@
+constant uint iq3_s_grid[512] = {
+    0x01010101, 0x01010103, 0x01010105, 0x0101010b, 0x0101010f, 0x01010301, 0x01010303, 0x01010305,
+    0x01010309, 0x0101030d, 0x01010501, 0x01010503, 0x0101050b, 0x01010707, 0x01010901, 0x01010905,
+    0x0101090b, 0x0101090f, 0x01010b03, 0x01010b07, 0x01010d01, 0x01010d05, 0x01010f03, 0x01010f09,
+    0x01010f0f, 0x01030101, 0x01030103, 0x01030105, 0x01030109, 0x01030301, 0x01030303, 0x0103030b,
+    0x01030501, 0x01030507, 0x0103050f, 0x01030703, 0x0103070b, 0x01030909, 0x01030d03, 0x01030d0b,
+    0x01030f05, 0x01050101, 0x01050103, 0x0105010b, 0x0105010f, 0x01050301, 0x01050307, 0x0105030d,
+    0x01050503, 0x0105050b, 0x01050701, 0x01050709, 0x01050905, 0x0105090b, 0x0105090f, 0x01050b03,
+    0x01050b07, 0x01050f01, 0x01050f07, 0x01070107, 0x01070303, 0x0107030b, 0x01070501, 0x01070505,
+    0x01070703, 0x01070707, 0x0107070d, 0x01070909, 0x01070b01, 0x01070b05, 0x01070d0f, 0x01070f03,
+    0x01070f0b, 0x01090101, 0x01090307, 0x0109030f, 0x01090503, 0x01090509, 0x01090705, 0x01090901,
+    0x01090907, 0x01090b03, 0x01090f01, 0x010b0105, 0x010b0109, 0x010b0501, 0x010b0505, 0x010b050d,
+    0x010b0707, 0x010b0903, 0x010b090b, 0x010b090f, 0x010b0d0d, 0x010b0f07, 0x010d010d, 0x010d0303,
+    0x010d0307, 0x010d0703, 0x010d0b05, 0x010d0f03, 0x010f0101, 0x010f0105, 0x010f0109, 0x010f0501,
+    0x010f0505, 0x010f050d, 0x010f0707, 0x010f0b01, 0x010f0b09, 0x03010101, 0x03010103, 0x03010105,
+    0x03010109, 0x03010301, 0x03010303, 0x03010307, 0x0301030b, 0x0301030f, 0x03010501, 0x03010505,
+    0x03010703, 0x03010709, 0x0301070d, 0x03010b09, 0x03010b0d, 0x03010d03, 0x03010f05, 0x03030101,
+    0x03030103, 0x03030107, 0x0303010d, 0x03030301, 0x03030309, 0x03030503, 0x03030701, 0x03030707,
+    0x03030903, 0x03030b01, 0x03030b05, 0x03030f01, 0x03030f0d, 0x03050101, 0x03050305, 0x0305030b,
+    0x0305030f, 0x03050501, 0x03050509, 0x03050705, 0x03050901, 0x03050907, 0x03050b0b, 0x03050d01,
+    0x03050f05, 0x03070103, 0x03070109, 0x0307010f, 0x03070301, 0x03070307, 0x03070503, 0x0307050f,
+    0x03070701, 0x03070709, 0x03070903, 0x03070d05, 0x03070f01, 0x03090107, 0x0309010b, 0x03090305,
+    0x03090309, 0x03090703, 0x03090707, 0x03090905, 0x0309090d, 0x03090b01, 0x03090b09, 0x030b0103,
+    0x030b0301, 0x030b0307, 0x030b0503, 0x030b0701, 0x030b0705, 0x030b0b03, 0x030d0501, 0x030d0509,
+    0x030d050f, 0x030d0909, 0x030d090d, 0x030f0103, 0x030f0107, 0x030f0301, 0x030f0305, 0x030f0503,
+    0x030f070b, 0x030f0903, 0x030f0d05, 0x030f0f01, 0x05010101, 0x05010103, 0x05010107, 0x0501010b,
+    0x0501010f, 0x05010301, 0x05010305, 0x05010309, 0x0501030d, 0x05010503, 0x05010507, 0x0501050f,
+    0x05010701, 0x05010705, 0x05010903, 0x05010907, 0x0501090b, 0x05010b01, 0x05010b05, 0x05010d0f,
+    0x05010f01, 0x05010f07, 0x05010f0b, 0x05030101, 0x05030105, 0x05030301, 0x05030307, 0x0503030f,
+    0x05030505, 0x0503050b, 0x05030703, 0x05030709, 0x05030905, 0x05030b03, 0x05050103, 0x05050109,
+    0x0505010f, 0x05050503, 0x05050507, 0x05050701, 0x0505070f, 0x05050903, 0x05050b07, 0x05050b0f,
+    0x05050f03, 0x05050f09, 0x05070101, 0x05070105, 0x0507010b, 0x05070303, 0x05070505, 0x05070509,
+    0x05070703, 0x05070707, 0x05070905, 0x05070b01, 0x05070d0d, 0x05090103, 0x0509010f, 0x05090501,
+    0x05090507, 0x05090705, 0x0509070b, 0x05090903, 0x05090f05, 0x05090f0b, 0x050b0109, 0x050b0303,
+    0x050b0505, 0x050b070f, 0x050b0901, 0x050b0b07, 0x050b0f01, 0x050d0101, 0x050d0105, 0x050d010f,
+    0x050d0503, 0x050d0b0b, 0x050d0d03, 0x050f010b, 0x050f0303, 0x050f050d, 0x050f0701, 0x050f0907,
+    0x050f0b01, 0x07010105, 0x07010303, 0x07010307, 0x0701030b, 0x0701030f, 0x07010505, 0x07010703,
+    0x07010707, 0x0701070b, 0x07010905, 0x07010909, 0x0701090f, 0x07010b03, 0x07010d07, 0x07010f03,
+    0x07030103, 0x07030107, 0x0703010b, 0x07030309, 0x07030503, 0x07030507, 0x07030901, 0x07030d01,
+    0x07030f05, 0x07030f0d, 0x07050101, 0x07050305, 0x07050501, 0x07050705, 0x07050709, 0x07050b01,
+    0x07070103, 0x07070301, 0x07070309, 0x07070503, 0x07070507, 0x0707050f, 0x07070701, 0x07070903,
+    0x07070907, 0x0707090f, 0x07070b0b, 0x07070f07, 0x07090107, 0x07090303, 0x0709030d, 0x07090505,
+    0x07090703, 0x07090b05, 0x07090d01, 0x07090d09, 0x070b0103, 0x070b0301, 0x070b0305, 0x070b050b,
+    0x070b0705, 0x070b0909, 0x070b0b0d, 0x070b0f07, 0x070d030d, 0x070d0903, 0x070f0103, 0x070f0107,
+    0x070f0501, 0x070f0505, 0x070f070b, 0x09010101, 0x09010109, 0x09010305, 0x09010501, 0x09010509,
+    0x0901050f, 0x09010705, 0x09010903, 0x09010b01, 0x09010f01, 0x09030105, 0x0903010f, 0x09030303,
+    0x09030307, 0x09030505, 0x09030701, 0x0903070b, 0x09030907, 0x09030b03, 0x09030b0b, 0x09050103,
+    0x09050107, 0x09050301, 0x0905030b, 0x09050503, 0x09050707, 0x09050901, 0x09050b0f, 0x09050d05,
+    0x09050f01, 0x09070109, 0x09070303, 0x09070307, 0x09070501, 0x09070505, 0x09070703, 0x0907070b,
+    0x09090101, 0x09090105, 0x09090509, 0x0909070f, 0x09090901, 0x09090f03, 0x090b010b, 0x090b010f,
+    0x090b0503, 0x090b0d05, 0x090d0307, 0x090d0709, 0x090d0d01, 0x090f0301, 0x090f030b, 0x090f0701,
+    0x090f0907, 0x090f0b03, 0x0b010105, 0x0b010301, 0x0b010309, 0x0b010505, 0x0b010901, 0x0b010909,
+    0x0b01090f, 0x0b010b05, 0x0b010d0d, 0x0b010f09, 0x0b030103, 0x0b030107, 0x0b03010b, 0x0b030305,
+    0x0b030503, 0x0b030705, 0x0b030f05, 0x0b050101, 0x0b050303, 0x0b050507, 0x0b050701, 0x0b05070d,
+    0x0b050b07, 0x0b070105, 0x0b07010f, 0x0b070301, 0x0b07050f, 0x0b070909, 0x0b070b03, 0x0b070d0b,
+    0x0b070f07, 0x0b090103, 0x0b090109, 0x0b090501, 0x0b090705, 0x0b09090d, 0x0b0b0305, 0x0b0b050d,
+    0x0b0b0b03, 0x0b0b0b07, 0x0b0d0905, 0x0b0f0105, 0x0b0f0109, 0x0b0f0505, 0x0d010303, 0x0d010307,
+    0x0d01030b, 0x0d010703, 0x0d010707, 0x0d010d01, 0x0d030101, 0x0d030501, 0x0d03050f, 0x0d030d09,
+    0x0d050305, 0x0d050709, 0x0d050905, 0x0d050b0b, 0x0d050d05, 0x0d050f01, 0x0d070101, 0x0d070309,
+    0x0d070503, 0x0d070901, 0x0d09050b, 0x0d090907, 0x0d090d05, 0x0d0b0101, 0x0d0b0107, 0x0d0b0709,
+    0x0d0b0d01, 0x0d0d010b, 0x0d0d0901, 0x0d0f0303, 0x0d0f0307, 0x0f010101, 0x0f010109, 0x0f01010f,
+    0x0f010501, 0x0f010505, 0x0f01070d, 0x0f010901, 0x0f010b09, 0x0f010d05, 0x0f030105, 0x0f030303,
+    0x0f030509, 0x0f030907, 0x0f03090b, 0x0f050103, 0x0f050109, 0x0f050301, 0x0f05030d, 0x0f050503,
+    0x0f050701, 0x0f050b03, 0x0f070105, 0x0f070705, 0x0f07070b, 0x0f070b07, 0x0f090103, 0x0f09010b,
+    0x0f090307, 0x0f090501, 0x0f090b01, 0x0f0b0505, 0x0f0b0905, 0x0f0d0105, 0x0f0d0703, 0x0f0f0101
+};

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

@@ -0,0 +1,38 @@
+constant uint iq3_xxs_grid[256] = {
+    0x04040404, 0x04040414, 0x04040424, 0x04040c0c, 0x04040c1c, 0x04040c3e, 0x04041404, 0x04041414,
+    0x04041c0c, 0x04042414, 0x04043e1c, 0x04043e2c, 0x040c040c, 0x040c041c, 0x040c0c04, 0x040c0c14,
+    0x040c140c, 0x040c142c, 0x040c1c04, 0x040c1c14, 0x040c240c, 0x040c2c24, 0x040c3e04, 0x04140404,
+    0x04140414, 0x04140424, 0x04140c0c, 0x04141404, 0x04141414, 0x04141c0c, 0x04141c1c, 0x04141c3e,
+    0x04142c0c, 0x04142c3e, 0x04143e2c, 0x041c040c, 0x041c043e, 0x041c0c04, 0x041c0c14, 0x041c142c,
+    0x041c3e04, 0x04240c1c, 0x04241c3e, 0x04242424, 0x04242c3e, 0x04243e1c, 0x04243e2c, 0x042c040c,
+    0x042c043e, 0x042c1c14, 0x042c2c14, 0x04341c2c, 0x04343424, 0x043e0c04, 0x043e0c24, 0x043e0c34,
+    0x043e241c, 0x043e340c, 0x0c04040c, 0x0c04041c, 0x0c040c04, 0x0c040c14, 0x0c04140c, 0x0c04141c,
+    0x0c041c04, 0x0c041c14, 0x0c041c24, 0x0c04243e, 0x0c042c04, 0x0c0c0404, 0x0c0c0414, 0x0c0c0c0c,
+    0x0c0c1404, 0x0c0c1414, 0x0c14040c, 0x0c14041c, 0x0c140c04, 0x0c140c14, 0x0c14140c, 0x0c141c04,
+    0x0c143e14, 0x0c1c0404, 0x0c1c0414, 0x0c1c1404, 0x0c1c1c0c, 0x0c1c2434, 0x0c1c3434, 0x0c24040c,
+    0x0c24042c, 0x0c242c04, 0x0c2c1404, 0x0c2c1424, 0x0c2c2434, 0x0c2c3e0c, 0x0c34042c, 0x0c3e1414,
+    0x0c3e2404, 0x14040404, 0x14040414, 0x14040c0c, 0x14040c1c, 0x14041404, 0x14041414, 0x14041434,
+    0x14041c0c, 0x14042414, 0x140c040c, 0x140c041c, 0x140c042c, 0x140c0c04, 0x140c0c14, 0x140c140c,
+    0x140c1c04, 0x140c341c, 0x140c343e, 0x140c3e04, 0x14140404, 0x14140414, 0x14140c0c, 0x14140c3e,
+    0x14141404, 0x14141414, 0x14141c3e, 0x14142404, 0x14142c2c, 0x141c040c, 0x141c0c04, 0x141c0c24,
+    0x141c3e04, 0x141c3e24, 0x14241c2c, 0x14242c1c, 0x142c041c, 0x142c143e, 0x142c240c, 0x142c3e24,
+    0x143e040c, 0x143e041c, 0x143e0c34, 0x143e242c, 0x1c04040c, 0x1c040c04, 0x1c040c14, 0x1c04140c,
+    0x1c04141c, 0x1c042c04, 0x1c04342c, 0x1c043e14, 0x1c0c0404, 0x1c0c0414, 0x1c0c1404, 0x1c0c1c0c,
+    0x1c0c2424, 0x1c0c2434, 0x1c14040c, 0x1c14041c, 0x1c140c04, 0x1c14142c, 0x1c142c14, 0x1c143e14,
+    0x1c1c0c0c, 0x1c1c1c1c, 0x1c241c04, 0x1c24243e, 0x1c243e14, 0x1c2c0404, 0x1c2c0434, 0x1c2c1414,
+    0x1c2c2c2c, 0x1c340c24, 0x1c341c34, 0x1c34341c, 0x1c3e1c1c, 0x1c3e3404, 0x24040424, 0x24040c3e,
+    0x24041c2c, 0x24041c3e, 0x24042c1c, 0x24042c3e, 0x240c3e24, 0x24141404, 0x24141c3e, 0x24142404,
+    0x24143404, 0x24143434, 0x241c043e, 0x241c242c, 0x24240424, 0x24242c0c, 0x24243424, 0x242c142c,
+    0x242c241c, 0x242c3e04, 0x243e042c, 0x243e0c04, 0x243e0c14, 0x243e1c04, 0x2c040c14, 0x2c04240c,
+    0x2c043e04, 0x2c0c0404, 0x2c0c0434, 0x2c0c1434, 0x2c0c2c2c, 0x2c140c24, 0x2c141c14, 0x2c143e14,
+    0x2c1c0414, 0x2c1c2c1c, 0x2c240c04, 0x2c24141c, 0x2c24143e, 0x2c243e14, 0x2c2c0414, 0x2c2c1c0c,
+    0x2c342c04, 0x2c3e1424, 0x2c3e2414, 0x34041424, 0x34042424, 0x34042434, 0x34043424, 0x340c140c,
+    0x340c340c, 0x34140c3e, 0x34143424, 0x341c1c04, 0x341c1c34, 0x34242424, 0x342c042c, 0x342c2c14,
+    0x34341c1c, 0x343e041c, 0x343e140c, 0x3e04041c, 0x3e04042c, 0x3e04043e, 0x3e040c04, 0x3e041c14,
+    0x3e042c14, 0x3e0c1434, 0x3e0c2404, 0x3e140c14, 0x3e14242c, 0x3e142c14, 0x3e1c0404, 0x3e1c0c2c,
+    0x3e1c1c1c, 0x3e1c3404, 0x3e24140c, 0x3e24240c, 0x3e2c0404, 0x3e2c0414, 0x3e2c1424, 0x3e341c04
+};
+
+inline uint iq3_xxs_signs(uint code) {
+    return code | ((popcount(code) & 1u) << 7);
+}

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

@@ -0,0 +1,10 @@
+constant char iq4_xs_values[16] = {
+    -127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113,
+};
+
+inline uint iq4_xs_code(uint n) {
+    const uint lo = (n & 4u) ? 0xF6EADDCFu : 0xBFAD9881u;
+    const uint hi = (n & 4u) ? 0x71594535u : 0x26190D01u;
+    const uint t = (n & 8u) ? hi : lo;
+    return (t >> (8*(n & 3u))) & 255u;
+}

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

@@ -0,0 +1,193 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_intel_required_subgroup_size
+#define INTEL_GPU 1
+#endif
+
+#define QK_K 256
+
+struct block_iq3_s {
+    half d;
+    uchar qs[QK_K/4];
+    uchar qh[QK_K/32];
+    uchar signs[QK_K/8];
+    uchar scales[QK_K/64];
+};
+
+#define LOAD_VEC_A 4
+#define LOAD_VEC_B 4
+
+#define BM 64
+#define BN 64
+#define BK 32
+#ifdef INTEL_GPU
+#define TM 8
+#define TN 8
+#else
+#define TM 4
+#define TN 8
+#endif
+
+kernel void kernel_mul_mm_iq3_s_f32_l4_lm(
+    global uchar * src0_qs,
+    global uchar * src0_sg,
+    global uchar * src0_qh,
+    global uchar * src0_sc,
+    global half * src0_d,
+    ulong offset0,
+    global float4 * src1,
+    ulong offset1,
+    global float  * dst,
+    ulong offsetd,
+
+    int ne00,
+    int ne01,
+    int ne02,
+    int ne11,
+    int ne12,
+
+    int stride_a,
+    int stride_b,
+    int stride_d,
+
+    int batch_stride_a,
+    int batch_stride_b,
+    int batch_stride_d,
+
+    int r2,
+    int r3
+) {
+    src1 = (global float4*)((global char*)src1 + offset1);
+    dst  = (global float *)((global char*)dst  + offsetd);
+
+    local float buf_a[BM * BK];
+    local float buf_b[BN * BK];
+
+    const int batch_idx = get_global_id(2);
+
+    const int i13 = batch_idx / ne12;
+    const int i12 = batch_idx % ne12;
+
+    const int i03 = i13 / r3;
+    const int i02 = i12 / r2;
+
+    const int batch_idx_a = i03 * ne02 + i02;
+
+    const int ir = get_group_id(0);
+    const int ic = get_group_id(1);
+
+    const int tid = get_local_id(0);
+    const int th_r  = tid % (BM / TM);
+    const int th_c  = tid / (BM / TM);
+
+    const int loadr_a = get_local_id(0) % (BK / LOAD_VEC_A);
+    const int loadc_a = get_local_id(0) / (BK / LOAD_VEC_A);
+    const int loadr_b = get_local_id(0) % (BK / LOAD_VEC_B);
+    const int loadc_b = get_local_id(0) / (BK / LOAD_VEC_B);
+
+    const int loadstride_a = get_local_size(0) * LOAD_VEC_A / BK;
+    const int loadstride_b = get_local_size(0) * LOAD_VEC_B / BK;
+
+    // pos_a counts elements, not blocks - one iq3_s super block holds QK_K of them
+    int pos_a = batch_idx_a * batch_stride_a + ir * BM * stride_a;
+    int pos_b = (batch_idx   * batch_stride_b + ic * BN * stride_b) / LOAD_VEC_B;
+
+    float sums[TM * TN];
+    float cache_a[TM];
+    float cache_b[TN];
+
+    for (int i = 0; i < TM * TN; i++) {
+        sums[i] = 0.0f;
+    }
+
+    for (int block = 0; block < ne00; block += BK) {
+        for (int l = 0; l < BM; l += loadstride_a) {
+            if (ir*BM + loadc_a + l < ne01) {
+                int idx = pos_a + (loadc_a + l) * stride_a + loadr_a * LOAD_VEC_A;
+                int ib  = idx / QK_K;
+                ib += offset0/sizeof(struct block_iq3_s);
+                int e   = idx % QK_K;
+
+                int ib32 = e >> 5;
+                int rem  = e & 31;
+                int lg    = rem >> 3;
+                int j0   = rem & 7;
+
+                uchar scb = src0_sc[4*ib + ib32/2];
+                float db = (float)src0_d[ib] * (float)(1 + 2*((scb >> (4*(ib32 & 1))) & 15));
+                int which = j0 >> 2;
+                uint gi = (uint)src0_qs[64*ib + 8*ib32 + 2*lg + which]
+                        | (((uint)src0_qh[8*ib + ib32] << (8 - 2*lg - which)) & 256);
+                uchar sg = src0_sg[32*ib + 4*ib32 + lg];
+                uint g = iq3_s_grid[gi];
+
+                float4 v1;
+                v1.s0 = db * (float)((g >>  0) & 0xFF) * ((sg & (1 << (j0+0))) ? -1.f : 1.f);
+                v1.s1 = db * (float)((g >>  8) & 0xFF) * ((sg & (1 << (j0+1))) ? -1.f : 1.f);
+                v1.s2 = db * (float)((g >> 16) & 0xFF) * ((sg & (1 << (j0+2))) ? -1.f : 1.f);
+                v1.s3 = db * (float)((g >> 24) & 0xFF) * ((sg & (1 << (j0+3))) ? -1.f : 1.f);
+
+                buf_a[(loadr_a * LOAD_VEC_A + 0) * BM + loadc_a + l] = v1.s0;
+                buf_a[(loadr_a * LOAD_VEC_A + 1) * BM + loadc_a + l] = v1.s1;
+                buf_a[(loadr_a * LOAD_VEC_A + 2) * BM + loadc_a + l] = v1.s2;
+                buf_a[(loadr_a * LOAD_VEC_A + 3) * BM + loadc_a + l] = v1.s3;
+            } else {
+                buf_a[(loadr_a * LOAD_VEC_A + 0) * BM + loadc_a + l] = 0.0f;
+                buf_a[(loadr_a * LOAD_VEC_A + 1) * BM + loadc_a + l] = 0.0f;
+                buf_a[(loadr_a * LOAD_VEC_A + 2) * BM + loadc_a + l] = 0.0f;
+                buf_a[(loadr_a * LOAD_VEC_A + 3) * BM + loadc_a + l] = 0.0f;
+            }
+        }
+
+        for (int l = 0; l < BN; l += loadstride_b) {
+            if (ic*BN + loadc_b + l < ne11) {
+                int idx = pos_b + (loadc_b + l) * stride_b / LOAD_VEC_B + loadr_b;
+                buf_b[(loadr_b * LOAD_VEC_B + 0) * BN + loadc_b + l] = src1[idx].s0;
+                buf_b[(loadr_b * LOAD_VEC_B + 1) * BN + loadc_b + l] = src1[idx].s1;
+                buf_b[(loadr_b * LOAD_VEC_B + 2) * BN + loadc_b + l] = src1[idx].s2;
+                buf_b[(loadr_b * LOAD_VEC_B + 3) * BN + loadc_b + l] = src1[idx].s3;
+            } else {
+                buf_b[(loadr_b * LOAD_VEC_B + 0) * BN + loadc_b + l] = 0.0f;
+                buf_b[(loadr_b * LOAD_VEC_B + 1) * BN + loadc_b + l] = 0.0f;
+                buf_b[(loadr_b * LOAD_VEC_B + 2) * BN + loadc_b + l] = 0.0f;
+                buf_b[(loadr_b * LOAD_VEC_B + 3) * BN + loadc_b + l] = 0.0f;
+            }
+        }
+
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+        pos_a += BK;
+        pos_b += BK / LOAD_VEC_B;
+
+        for (int i = 0; i < BK; i++) {
+            for (int j = 0; j < TM; j++) {
+                cache_a[j] = buf_a[(i) * BM + th_r * TM + j];
+            }
+
+            for (int j = 0; j < TN; j++) {
+                cache_b[j] = buf_b[(i) * BN + th_c * TN + j];
+            }
+
+            for (int cc = 0; cc < TN; cc++) {
+                for (int cr = 0; cr < TM; cr++) {
+                    const int sums_idx = cc*TM + cr;
+                    sums[sums_idx] = mad(cache_a[cr], cache_b[cc], sums[sums_idx]);
+                }
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    const int dr = ir * BM + th_r * TM;
+    const int dc = ic * BN + th_c * TN;
+
+    const ulong offsets = batch_idx * batch_stride_d;
+
+    for (int cc = 0; cc < TN; cc++) {
+        for (int cr = 0; cr < TM; cr++) {
+            if (dr + cr < ne01 && dc + cc < ne11) {
+                dst[offsets + (dc + cc) * stride_d + dr + cr] = sums[cc * TM + cr];
+            }
+        }
+    }
+}

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

@@ -0,0 +1,187 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_intel_required_subgroup_size
+#define INTEL_GPU 1
+#endif
+
+#define QK_K 256
+
+struct block_iq3_xxs {
+    half d;
+    uchar qs[3*QK_K/8];
+};
+
+#define LOAD_VEC_A 4
+#define LOAD_VEC_B 4
+
+#define BM 64
+#define BN 64
+#define BK 32
+#ifdef INTEL_GPU
+#define TM 8
+#define TN 8
+#else
+#define TM 4
+#define TN 8
+#endif
+
+kernel void kernel_mul_mm_iq3_xxs_f32_l4_lm(
+    global uchar * src0_qs,
+    global uint * src0_sas,
+    global half * src0_d,
+    ulong offset0,
+    global float4 * src1,
+    ulong offset1,
+    global float  * dst,
+    ulong offsetd,
+
+    int ne00,
+    int ne01,
+    int ne02,
+    int ne11,
+    int ne12,
+
+    int stride_a,
+    int stride_b,
+    int stride_d,
+
+    int batch_stride_a,
+    int batch_stride_b,
+    int batch_stride_d,
+
+    int r2,
+    int r3
+) {
+    src1 = (global float4*)((global char*)src1 + offset1);
+    dst  = (global float *)((global char*)dst  + offsetd);
+
+    local float buf_a[BM * BK];
+    local float buf_b[BN * BK];
+
+    const int batch_idx = get_global_id(2);
+
+    const int i13 = batch_idx / ne12;
+    const int i12 = batch_idx % ne12;
+
+    const int i03 = i13 / r3;
+    const int i02 = i12 / r2;
+
+    const int batch_idx_a = i03 * ne02 + i02;
+
+    const int ir = get_group_id(0);
+    const int ic = get_group_id(1);
+
+    const int tid = get_local_id(0);
+    const int th_r  = tid % (BM / TM);
+    const int th_c  = tid / (BM / TM);
+
+    const int loadr_a = get_local_id(0) % (BK / LOAD_VEC_A);
+    const int loadc_a = get_local_id(0) / (BK / LOAD_VEC_A);
+    const int loadr_b = get_local_id(0) % (BK / LOAD_VEC_B);
+    const int loadc_b = get_local_id(0) / (BK / LOAD_VEC_B);
+
+    const int loadstride_a = get_local_size(0) * LOAD_VEC_A / BK;
+    const int loadstride_b = get_local_size(0) * LOAD_VEC_B / BK;
+
+    // pos_a counts elements, not blocks - one iq3_xxs super block holds QK_K of them
+    int pos_a = batch_idx_a * batch_stride_a + ir * BM * stride_a;
+    int pos_b = (batch_idx   * batch_stride_b + ic * BN * stride_b) / LOAD_VEC_B;
+
+    float sums[TM * TN];
+    float cache_a[TM];
+    float cache_b[TN];
+
+    for (int i = 0; i < TM * TN; i++) {
+        sums[i] = 0.0f;
+    }
+
+    for (int block = 0; block < ne00; block += BK) {
+        for (int l = 0; l < BM; l += loadstride_a) {
+            if (ir*BM + loadc_a + l < ne01) {
+                int idx = pos_a + (loadc_a + l) * stride_a + loadr_a * LOAD_VEC_A;
+                int ib  = idx / QK_K;
+                ib += offset0/sizeof(struct block_iq3_xxs);
+                int e   = idx % QK_K;
+
+                int ib32 = e >> 5;
+                int rem  = e & 31;
+                int lg    = rem >> 3;
+                int j0   = rem & 7;
+
+                uint aux = src0_sas[8*ib + ib32];
+                float db = (float)src0_d[ib] * (0.5f + (float)(aux >> 28)) * 0.5f;
+
+                uint gi = src0_qs[64*ib + 8*ib32 + 2*lg + (j0 >> 2)];
+                uint sg = iq3_xxs_signs((aux >> (7*lg)) & 127u);
+                uint  g  = iq3_xxs_grid[gi];
+
+                float4 v1;
+                v1.s0 = db * (float)((g >>  0) & 0xFF) * ((sg & (1 << (j0+0))) ? -1.f : 1.f);
+                v1.s1 = db * (float)((g >>  8) & 0xFF) * ((sg & (1 << (j0+1))) ? -1.f : 1.f);
+                v1.s2 = db * (float)((g >> 16) & 0xFF) * ((sg & (1 << (j0+2))) ? -1.f : 1.f);
+                v1.s3 = db * (float)((g >> 24) & 0xFF) * ((sg & (1 << (j0+3))) ? -1.f : 1.f);
+
+                buf_a[(loadr_a * LOAD_VEC_A + 0) * BM + loadc_a + l] = v1.s0;
+                buf_a[(loadr_a * LOAD_VEC_A + 1) * BM + loadc_a + l] = v1.s1;
+                buf_a[(loadr_a * LOAD_VEC_A + 2) * BM + loadc_a + l] = v1.s2;
+                buf_a[(loadr_a * LOAD_VEC_A + 3) * BM + loadc_a + l] = v1.s3;
+            } else {
+                buf_a[(loadr_a * LOAD_VEC_A + 0) * BM + loadc_a + l] = 0.0f;
+                buf_a[(loadr_a * LOAD_VEC_A + 1) * BM + loadc_a + l] = 0.0f;
+                buf_a[(loadr_a * LOAD_VEC_A + 2) * BM + loadc_a + l] = 0.0f;
+                buf_a[(loadr_a * LOAD_VEC_A + 3) * BM + loadc_a + l] = 0.0f;
+            }
+        }
+
+        for (int l = 0; l < BN; l += loadstride_b) {
+            if (ic*BN + loadc_b + l < ne11) {
+                int idx = pos_b + (loadc_b + l) * stride_b / LOAD_VEC_B + loadr_b;
+                buf_b[(loadr_b * LOAD_VEC_B + 0) * BN + loadc_b + l] = src1[idx].s0;
+                buf_b[(loadr_b * LOAD_VEC_B + 1) * BN + loadc_b + l] = src1[idx].s1;
+                buf_b[(loadr_b * LOAD_VEC_B + 2) * BN + loadc_b + l] = src1[idx].s2;
+                buf_b[(loadr_b * LOAD_VEC_B + 3) * BN + loadc_b + l] = src1[idx].s3;
+            } else {
+                buf_b[(loadr_b * LOAD_VEC_B + 0) * BN + loadc_b + l] = 0.0f;
+                buf_b[(loadr_b * LOAD_VEC_B + 1) * BN + loadc_b + l] = 0.0f;
+                buf_b[(loadr_b * LOAD_VEC_B + 2) * BN + loadc_b + l] = 0.0f;
+                buf_b[(loadr_b * LOAD_VEC_B + 3) * BN + loadc_b + l] = 0.0f;
+            }
+        }
+
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+        pos_a += BK;
+        pos_b += BK / LOAD_VEC_B;
+
+        for (int i = 0; i < BK; i++) {
+            for (int j = 0; j < TM; j++) {
+                cache_a[j] = buf_a[(i) * BM + th_r * TM + j];
+            }
+
+            for (int j = 0; j < TN; j++) {
+                cache_b[j] = buf_b[(i) * BN + th_c * TN + j];
+            }
+
+            for (int cc = 0; cc < TN; cc++) {
+                for (int cr = 0; cr < TM; cr++) {
+                    const int sums_idx = cc*TM + cr;
+                    sums[sums_idx] = mad(cache_a[cr], cache_b[cc], sums[sums_idx]);
+                }
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    const int dr = ir * BM + th_r * TM;
+    const int dc = ic * BN + th_c * TN;
+
+    const ulong offsets = batch_idx * batch_stride_d;
+
+    for (int cc = 0; cc < TN; cc++) {
+        for (int cr = 0; cr < TM; cr++) {
+            if (dr + cr < ne01 && dc + cc < ne11) {
+                dst[offsets + (dc + cc) * stride_d + dr + cr] = sums[cc * TM + cr];
+            }
+        }
+    }
+}

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

@@ -0,0 +1,186 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_intel_required_subgroup_size
+#define INTEL_GPU 1
+#endif
+
+#define QK_K 256
+
+struct block_iq4_xs {
+    half d;
+    ushort scales_h;
+    uchar scales_l[QK_K/64];
+    uchar qs[QK_K/2];
+};
+
+#define LOAD_VEC_A 4
+#define LOAD_VEC_B 4
+
+#define BM 64
+#define BN 64
+#define BK 32
+#ifdef INTEL_GPU
+#define TM 8
+#define TN 8
+#else
+#define TM 4
+#define TN 8
+#endif
+
+kernel void kernel_mul_mm_iq4_xs_f32_l4_lm(
+    global uchar * src0_qs,
+    global uchar * src0_sl,
+    global ushort * src0_sh,
+    global half * src0_d,
+    ulong offset0,
+    global float4 * src1,
+    ulong offset1,
+    global float  * dst,
+    ulong offsetd,
+
+    int ne00,
+    int ne01,
+    int ne02,
+    int ne11,
+    int ne12,
+
+    int stride_a,
+    int stride_b,
+    int stride_d,
+
+    int batch_stride_a,
+    int batch_stride_b,
+    int batch_stride_d,
+
+    int r2,
+    int r3
+) {
+    src1 = (global float4*)((global char*)src1 + offset1);
+    dst  = (global float *)((global char*)dst  + offsetd);
+
+    local float buf_a[BM * BK];
+    local float buf_b[BN * BK];
+
+    const int batch_idx = get_global_id(2);
+
+    const int i13 = batch_idx / ne12;
+    const int i12 = batch_idx % ne12;
+
+    const int i03 = i13 / r3;
+    const int i02 = i12 / r2;
+
+    const int batch_idx_a = i03 * ne02 + i02;
+
+    const int ir = get_group_id(0);
+    const int ic = get_group_id(1);
+
+    const int tid = get_local_id(0);
+    const int th_r  = tid % (BM / TM);
+    const int th_c  = tid / (BM / TM);
+
+    const int loadr_a = get_local_id(0) % (BK / LOAD_VEC_A);
+    const int loadc_a = get_local_id(0) / (BK / LOAD_VEC_A);
+    const int loadr_b = get_local_id(0) % (BK / LOAD_VEC_B);
+    const int loadc_b = get_local_id(0) / (BK / LOAD_VEC_B);
+
+    const int loadstride_a = get_local_size(0) * LOAD_VEC_A / BK;
+    const int loadstride_b = get_local_size(0) * LOAD_VEC_B / BK;
+
+    // pos_a counts elements, not blocks - one iq4_xs super block holds QK_K of them
+    int pos_a = batch_idx_a * batch_stride_a + ir * BM * stride_a;
+    int pos_b = (batch_idx   * batch_stride_b + ic * BN * stride_b) / LOAD_VEC_B;
+
+    float sums[TM * TN];
+    float cache_a[TM];
+    float cache_b[TN];
+
+    for (int i = 0; i < TM * TN; i++) {
+        sums[i] = 0.0f;
+    }
+
+    for (int block = 0; block < ne00; block += BK) {
+        for (int l = 0; l < BM; l += loadstride_a) {
+            if (ir*BM + loadc_a + l < ne01) {
+                int idx = pos_a + (loadc_a + l) * stride_a + loadr_a * LOAD_VEC_A;
+                int ib  = idx / QK_K;
+                ib += offset0/sizeof(struct block_iq4_xs);
+                int e   = idx % QK_K;
+
+                int ib32 = e >> 5;
+                int j = e & 15;
+                int shift = (e & 16) >> 2;
+                int ls = ((src0_sl[4*ib + ib32/2] >> (4*(ib32 & 1))) & 15) | (((src0_sh[ib] >> (2*ib32)) & 3) << 4);
+                float db = (float)src0_d[ib] * (float)(ls - 32);
+                uchar4 q = vload4(0, src0_qs + 128*ib + 16*ib32 + j);
+                q = (q >> (uchar4)(shift)) & (uchar4)(15);
+
+                float4 v1;
+                v1.s0 = db * (float)iq4_xs_values[q.s0];
+                v1.s1 = db * (float)iq4_xs_values[q.s1];
+                v1.s2 = db * (float)iq4_xs_values[q.s2];
+                v1.s3 = db * (float)iq4_xs_values[q.s3];
+
+                buf_a[(loadr_a * LOAD_VEC_A + 0) * BM + loadc_a + l] = v1.s0;
+                buf_a[(loadr_a * LOAD_VEC_A + 1) * BM + loadc_a + l] = v1.s1;
+                buf_a[(loadr_a * LOAD_VEC_A + 2) * BM + loadc_a + l] = v1.s2;
+                buf_a[(loadr_a * LOAD_VEC_A + 3) * BM + loadc_a + l] = v1.s3;
+            } else {
+                buf_a[(loadr_a * LOAD_VEC_A + 0) * BM + loadc_a + l] = 0.0f;
+                buf_a[(loadr_a * LOAD_VEC_A + 1) * BM + loadc_a + l] = 0.0f;
+                buf_a[(loadr_a * LOAD_VEC_A + 2) * BM + loadc_a + l] = 0.0f;
+                buf_a[(loadr_a * LOAD_VEC_A + 3) * BM + loadc_a + l] = 0.0f;
+            }
+        }
+
+        for (int l = 0; l < BN; l += loadstride_b) {
+            if (ic*BN + loadc_b + l < ne11) {
+                int idx = pos_b + (loadc_b + l) * stride_b / LOAD_VEC_B + loadr_b;
+                buf_b[(loadr_b * LOAD_VEC_B + 0) * BN + loadc_b + l] = src1[idx].s0;
+                buf_b[(loadr_b * LOAD_VEC_B + 1) * BN + loadc_b + l] = src1[idx].s1;
+                buf_b[(loadr_b * LOAD_VEC_B + 2) * BN + loadc_b + l] = src1[idx].s2;
+                buf_b[(loadr_b * LOAD_VEC_B + 3) * BN + loadc_b + l] = src1[idx].s3;
+            } else {
+                buf_b[(loadr_b * LOAD_VEC_B + 0) * BN + loadc_b + l] = 0.0f;
+                buf_b[(loadr_b * LOAD_VEC_B + 1) * BN + loadc_b + l] = 0.0f;
+                buf_b[(loadr_b * LOAD_VEC_B + 2) * BN + loadc_b + l] = 0.0f;
+                buf_b[(loadr_b * LOAD_VEC_B + 3) * BN + loadc_b + l] = 0.0f;
+            }
+        }
+
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+        pos_a += BK;
+        pos_b += BK / LOAD_VEC_B;
+
+        for (int i = 0; i < BK; i++) {
+            for (int j = 0; j < TM; j++) {
+                cache_a[j] = buf_a[(i) * BM + th_r * TM + j];
+            }
+
+            for (int j = 0; j < TN; j++) {
+                cache_b[j] = buf_b[(i) * BN + th_c * TN + j];
+            }
+
+            for (int cc = 0; cc < TN; cc++) {
+                for (int cr = 0; cr < TM; cr++) {
+                    const int sums_idx = cc*TM + cr;
+                    sums[sums_idx] = mad(cache_a[cr], cache_b[cc], sums[sums_idx]);
+                }
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    const int dr = ir * BM + th_r * TM;
+    const int dc = ic * BN + th_c * TN;
+
+    const ulong offsets = batch_idx * batch_stride_d;
+
+    for (int cc = 0; cc < TN; cc++) {
+        for (int cr = 0; cr < TM; cr++) {
+            if (dr + cr < ne01 && dc + cc < ne11) {
+                dst[offsets + (dc + cc) * stride_d + dr + cr] = sums[cc * TM + cr];
+            }
+        }
+    }
+}

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

@@ -0,0 +1,146 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_intel_required_subgroup_size
+#pragma OPENCL EXTENSION cl_intel_required_subgroup_size : enable
+#define INTEL_GPU 1
+#define REQD_SUBGROUP_SIZE_16 __attribute__((intel_reqd_sub_group_size(16)))
+#define REQD_SUBGROUP_SIZE_32 __attribute__((intel_reqd_sub_group_size(32)))
+#elif defined(cl_qcom_reqd_sub_group_size)
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+#define ADRENO_GPU 1
+#define REQD_SUBGROUP_SIZE_64  __attribute__((qcom_reqd_sub_group_size("half")))
+#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+
+#define QK_K 256
+#define IQ3_S_N_SCALE (QK_K/64)
+
+typedef struct {
+    half  d;
+    uchar qs[QK_K/4];
+    uchar qh[QK_K/32];
+    uchar signs[QK_K/8];
+    uchar scales[IQ3_S_N_SCALE];
+} block_iq3_s;
+
+#undef N_DST
+#undef N_SIMDGROUP
+#undef N_SIMDWIDTH
+
+#ifdef INTEL_GPU
+#define N_DST 4
+#define N_SIMDGROUP 1
+#define N_SIMDWIDTH 16
+#elif defined (ADRENO_GPU)
+#define N_DST 4
+#define N_SIMDGROUP 1
+#define N_SIMDWIDTH 64
+#endif
+
+#undef  BLOCK_STRIDE
+// 8 threads cover one super block, one 32 element sub block each
+#define BLOCK_STRIDE (N_SIMDWIDTH/8)
+
+#ifdef INTEL_GPU
+REQD_SUBGROUP_SIZE_16
+#elif defined (ADRENO_GPU)
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mv_iq3_s_f32(
+        global char * src0,
+        ulong offset0,
+        global char * src1,
+        ulong offset1,
+        global char * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne12,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src0 = src0 + offset0;
+    src1 = src1 + offset1;
+    dst  = dst  + offsetd;
+
+    int ix = get_sub_group_local_id()/8;  // super block index
+    int it = get_sub_group_local_id()%8;  // sub block inside the super block
+
+    int nb = ne00/QK_K;
+
+    int r0 = get_group_id(0);
+    int r1 = get_group_id(1);
+    int im = get_group_id(2);
+    int first_row = (r0 * N_SIMDGROUP + get_sub_group_id()) * N_DST;
+
+    int i12 = im%ne12;
+    int i13 = im/ne12;
+
+    ulong offset_src0 = first_row*nb01 + (i12/r2)*nb02 + (i13/r3)*nb03;
+    ulong offset_src1 =        r1*nb11 + (i12   )*nb12 + (i13   )*nb13;
+
+    global block_iq3_s * x = (global block_iq3_s *) (src0 + offset_src0);
+    global float         * y = (global float         *) (src1 + offset_src1);
+
+    float yl[32];
+    float sumf[N_DST] = {0.f};
+    float all_sum;
+
+    global float * y4 = y + ix * QK_K + 32 * it;
+
+    for (int ib = ix; ib < nb; ib += BLOCK_STRIDE) {
+        for (int i = 0; i < 32; ++i) {
+            yl[i] = y4[i];
+        }
+
+        global char * xrow = (global char *)(x + ib);
+
+        // only ne01 output rows exist; reading past them can pull in an inf scale
+        for (int row = 0; row < N_DST && first_row + row < ne01; row++) {
+            global block_iq3_s * xb = (global block_iq3_s *)(xrow + row*nb01);
+
+            uchar nib = (it & 1) ? (xb->scales[it >> 1] >> 4) : (xb->scales[it >> 1] & 0xf);
+            float db  = (float)xb->d * (float)(1 + 2*nib);
+
+            uchar qhb = xb->qh[it];
+            global uchar * qsb = xb->qs + 8*it;
+            global uchar * sgb = xb->signs + 4*it;
+
+            float acc = 0.f;
+            for (int l = 0; l < 4; ++l) {
+                uint i1 = (uint)qsb[2*l+0] | (((uint)qhb << (8-2*l)) & 256);
+                uint i2 = (uint)qsb[2*l+1] | (((uint)qhb << (7-2*l)) & 256);
+                uint g1 = iq3_s_grid[i1];
+                uint g2 = iq3_s_grid[i2];
+                uchar sg = sgb[l];
+                for (int j = 0; j < 4; ++j) {
+                    acc += yl[8*l+j+0] * (float)((g1 >> (8*j)) & 0xFF) * ((sg & (1 << j))     ? -1.f : 1.f);
+                    acc += yl[8*l+j+4] * (float)((g2 >> (8*j)) & 0xFF) * ((sg & (1 << (j+4))) ? -1.f : 1.f);
+                }
+            }
+
+            sumf[row] += db * acc;
+        }
+
+        y4 += BLOCK_STRIDE * QK_K;
+    }
+
+    global float * dst_f32 = (global float *) dst + im*ne0*ne1 + r1*ne0;
+
+    for (int row = 0; row < N_DST; ++row) {
+        all_sum = sub_group_reduce_add(sumf[row]);
+        if (first_row + row < ne01) {
+            if (get_sub_group_local_id() == 0) {
+                dst_f32[first_row + row] = all_sum;
+            }
+        }
+    }
+}

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

@@ -0,0 +1,147 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_intel_required_subgroup_size
+#pragma OPENCL EXTENSION cl_intel_required_subgroup_size : enable
+#define INTEL_GPU 1
+#define REQD_SUBGROUP_SIZE_16 __attribute__((intel_reqd_sub_group_size(16)))
+#define REQD_SUBGROUP_SIZE_32 __attribute__((intel_reqd_sub_group_size(32)))
+#elif defined(cl_qcom_reqd_sub_group_size)
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+#define ADRENO_GPU 1
+#define REQD_SUBGROUP_SIZE_64  __attribute__((qcom_reqd_sub_group_size("half")))
+#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+
+#define QK_K 256
+#define IQ3_S_N_SCALE (QK_K/64)
+
+typedef struct {
+    half  d;
+    uchar qs[QK_K/4];
+    uchar qh[QK_K/32];
+    uchar signs[QK_K/8];
+    uchar scales[IQ3_S_N_SCALE];
+} block_iq3_s;
+
+#undef N_DST
+#undef N_SIMDGROUP
+#undef N_SIMDWIDTH
+
+#ifdef INTEL_GPU
+#define N_DST 4
+#define N_SIMDGROUP 1
+#define N_SIMDWIDTH 16
+#elif defined (ADRENO_GPU)
+#define N_DST 4
+#define N_SIMDGROUP 1
+#define N_SIMDWIDTH 64
+#endif
+
+#undef  BLOCK_STRIDE
+// 8 threads cover one super block, one 32 element sub block each
+#define BLOCK_STRIDE (N_SIMDWIDTH/8)
+
+#ifdef INTEL_GPU
+REQD_SUBGROUP_SIZE_16
+#elif defined (ADRENO_GPU)
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mv_iq3_s_f32_flat(
+        global uchar * src0_qs,
+        global uchar * src0_sg,
+        global uchar * src0_qh,
+        global uchar * src0_sc,
+        global half * src0_d,
+        ulong offset0,
+        global char * src1,
+        ulong offset1,
+        global char * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne12,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src1 = src1 + offset1;
+    dst  = dst  + offsetd;
+
+    int ix = get_sub_group_local_id()/8;  // super block index
+    int it = get_sub_group_local_id()%8;  // sub block inside the super block
+
+    int nb = ne00/QK_K;
+
+    int r0 = get_group_id(0);
+    int r1 = get_group_id(1);
+    int im = get_group_id(2);
+    int first_row = (r0 * N_SIMDGROUP + get_sub_group_id()) * N_DST;
+
+    int i12 = im%ne12;
+    int i13 = im/ne12;
+
+    ulong offset_src0 = (offset0 + first_row*nb01 + (i12/r2)*nb02 + (i13/r3)*nb03)/sizeof(block_iq3_s);
+    ulong row_stride = nb01/sizeof(block_iq3_s);
+    ulong offset_src1 =        r1*nb11 + (i12   )*nb12 + (i13   )*nb13;
+
+    global float         * y = (global float         *) (src1 + offset_src1);
+
+    float yl[32];
+    float sumf[N_DST] = {0.f};
+    float all_sum;
+
+    global float * y4 = y + ix * QK_K + 32 * it;
+
+    for (int ib = ix; ib < nb; ib += BLOCK_STRIDE) {
+        for (int i = 0; i < 32; ++i) {
+            yl[i] = y4[i];
+        }
+
+        // only ne01 output rows exist; reading past them can pull in an inf scale
+        for (int row = 0; row < N_DST && first_row + row < ne01; row++) {
+            const ulong b = offset_src0 + ib + row*row_stride;
+
+            uchar nib = (it & 1) ? (src0_sc[4*b + (it >> 1)] >> 4) : (src0_sc[4*b + (it >> 1)] & 0xf);
+            float db  = (float)src0_d[b] * (float)(1 + 2*nib);
+
+            uchar qhb = src0_qh[8*b + it];
+            global uchar * qsb = src0_qs + 64*b + 8*it;
+            global uchar * sgb = src0_sg + 32*b + 4*it;
+
+            float acc = 0.f;
+            for (int l = 0; l < 4; ++l) {
+                uint i1 = (uint)qsb[2*l+0] | (((uint)qhb << (8-2*l)) & 256);
+                uint i2 = (uint)qsb[2*l+1] | (((uint)qhb << (7-2*l)) & 256);
+                uint g1 = iq3_s_grid[i1];
+                uint g2 = iq3_s_grid[i2];
+                uchar sg = sgb[l];
+                for (int j = 0; j < 4; ++j) {
+                    acc += yl[8*l+j+0] * (float)((g1 >> (8*j)) & 0xFF) * ((sg & (1 << j))     ? -1.f : 1.f);
+                    acc += yl[8*l+j+4] * (float)((g2 >> (8*j)) & 0xFF) * ((sg & (1 << (j+4))) ? -1.f : 1.f);
+                }
+            }
+
+            sumf[row] += db * acc;
+        }
+
+        y4 += BLOCK_STRIDE * QK_K;
+    }
+
+    global float * dst_f32 = (global float *) dst + im*ne0*ne1 + r1*ne0;
+
+    for (int row = 0; row < N_DST; ++row) {
+        all_sum = sub_group_reduce_add(sumf[row]);
+        if (first_row + row < ne01) {
+            if (get_sub_group_local_id() == 0) {
+                dst_f32[first_row + row] = all_sum;
+            }
+        }
+    }
+}

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

@@ -0,0 +1,145 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_intel_required_subgroup_size
+#pragma OPENCL EXTENSION cl_intel_required_subgroup_size : enable
+#define INTEL_GPU 1
+#define REQD_SUBGROUP_SIZE_16 __attribute__((intel_reqd_sub_group_size(16)))
+#define REQD_SUBGROUP_SIZE_32 __attribute__((intel_reqd_sub_group_size(32)))
+#elif defined(cl_qcom_reqd_sub_group_size)
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+#define ADRENO_GPU 1
+#define REQD_SUBGROUP_SIZE_64  __attribute__((qcom_reqd_sub_group_size("half")))
+#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+
+#define QK_K 256
+
+// d, then 64 grid index bytes, then 8 uint32 of packed scale+signs
+typedef struct {
+    half  d;
+    uchar qs[3*QK_K/8];
+} block_iq3_xxs;
+
+#undef N_DST
+#undef N_SIMDGROUP
+#undef N_SIMDWIDTH
+
+#ifdef INTEL_GPU
+#define N_DST 4
+#define N_SIMDGROUP 1
+#define N_SIMDWIDTH 16
+#elif defined (ADRENO_GPU)
+#define N_DST 4
+#define N_SIMDGROUP 1
+#define N_SIMDWIDTH 64
+#endif
+
+#undef  BLOCK_STRIDE
+// 8 threads cover one super block, one 32 element sub block each
+#define BLOCK_STRIDE (N_SIMDWIDTH/8)
+
+#ifdef INTEL_GPU
+REQD_SUBGROUP_SIZE_16
+#elif defined (ADRENO_GPU)
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mv_iq3_xxs_f32(
+        global char * src0,
+        ulong offset0,
+        global char * src1,
+        ulong offset1,
+        global char * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne12,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src0 = src0 + offset0;
+    src1 = src1 + offset1;
+    dst  = dst  + offsetd;
+
+    int ix = get_sub_group_local_id()/8;  // super block index
+    int it = get_sub_group_local_id()%8;  // sub block inside the super block
+
+    int nb = ne00/QK_K;
+
+    int r0 = get_group_id(0);
+    int r1 = get_group_id(1);
+    int im = get_group_id(2);
+    int first_row = (r0 * N_SIMDGROUP + get_sub_group_id()) * N_DST;
+
+    int i12 = im%ne12;
+    int i13 = im/ne12;
+
+    ulong offset_src0 = first_row*nb01 + (i12/r2)*nb02 + (i13/r3)*nb03;
+    ulong offset_src1 =        r1*nb11 + (i12   )*nb12 + (i13   )*nb13;
+
+    global block_iq3_xxs * x = (global block_iq3_xxs *) (src0 + offset_src0);
+    global float         * y = (global float         *) (src1 + offset_src1);
+
+    float yl[32];
+    float sumf[N_DST] = {0.f};
+    float all_sum;
+
+    global float * y4 = y + ix * QK_K + 32 * it;
+
+    for (int ib = ix; ib < nb; ib += BLOCK_STRIDE) {
+        for (int i = 0; i < 32; ++i) {
+            yl[i] = y4[i];
+        }
+
+        global char * xrow = (global char *)(x + ib);
+
+        // only ne01 output rows exist; reading past them can pull in an inf scale
+        for (int row = 0; row < N_DST && first_row + row < ne01; row++) {
+            global block_iq3_xxs * xb = (global block_iq3_xxs *)(xrow + row*nb01);
+
+            global uchar * qs = xb->qs + 8*it;
+
+            // Blocks are 98 bytes, so packed signs are only 2-byte aligned.
+            global ushort * sw = (global ushort *)(xb->qs + 64 + 4*it);
+            uint aux32 = (uint)sw[0] | ((uint)sw[1] << 16);
+
+            float db = (float)xb->d * (0.5f + (float)(aux32 >> 28)) * 0.5f;
+
+            float acc = 0.f;
+            for (int l = 0; l < 4; ++l) {
+                uchar signs = iq3_xxs_signs((aux32 >> (7*l)) & 127u);
+                uint  g1 = iq3_xxs_grid[qs[2*l+0]];
+                uint  g2 = iq3_xxs_grid[qs[2*l+1]];
+
+                for (int j = 0; j < 4; ++j) {
+                    float s1 = (signs & (1 << j))       ? -1.f : 1.f;
+                    float s2 = (signs & (1 << (j+4)))   ? -1.f : 1.f;
+                    acc += yl[8*l+j+0] * (float)((g1 >> (8*j)) & 0xFF) * s1;
+                    acc += yl[8*l+j+4] * (float)((g2 >> (8*j)) & 0xFF) * s2;
+                }
+            }
+
+            sumf[row] += db * acc;
+        }
+
+        y4 += BLOCK_STRIDE * QK_K;
+    }
+
+    global float * dst_f32 = (global float *) dst + im*ne0*ne1 + r1*ne0;
+
+    for (int row = 0; row < N_DST; ++row) {
+        all_sum = sub_group_reduce_add(sumf[row]);
+        if (first_row + row < ne01) {
+            if (get_sub_group_local_id() == 0) {
+                dst_f32[first_row + row] = all_sum;
+            }
+        }
+    }
+}

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

@@ -0,0 +1,142 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_intel_required_subgroup_size
+#pragma OPENCL EXTENSION cl_intel_required_subgroup_size : enable
+#define INTEL_GPU 1
+#define REQD_SUBGROUP_SIZE_16 __attribute__((intel_reqd_sub_group_size(16)))
+#define REQD_SUBGROUP_SIZE_32 __attribute__((intel_reqd_sub_group_size(32)))
+#elif defined(cl_qcom_reqd_sub_group_size)
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+#define ADRENO_GPU 1
+#define REQD_SUBGROUP_SIZE_64  __attribute__((qcom_reqd_sub_group_size("half")))
+#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+
+#define QK_K 256
+
+// d, then 64 grid index bytes, then 8 uint32 of packed scale+signs
+typedef struct {
+    half  d;
+    uchar qs[3*QK_K/8];
+} block_iq3_xxs;
+
+#undef N_DST
+#undef N_SIMDGROUP
+#undef N_SIMDWIDTH
+
+#ifdef INTEL_GPU
+#define N_DST 4
+#define N_SIMDGROUP 1
+#define N_SIMDWIDTH 16
+#elif defined (ADRENO_GPU)
+#define N_DST 4
+#define N_SIMDGROUP 1
+#define N_SIMDWIDTH 64
+#endif
+
+#undef  BLOCK_STRIDE
+// 8 threads cover one super block, one 32 element sub block each
+#define BLOCK_STRIDE (N_SIMDWIDTH/8)
+
+#ifdef INTEL_GPU
+REQD_SUBGROUP_SIZE_16
+#elif defined (ADRENO_GPU)
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mv_iq3_xxs_f32_flat(
+        global uchar * src0_qs,
+        global uint * src0_sas,
+        global half * src0_d,
+        ulong offset0,
+        global char * src1,
+        ulong offset1,
+        global char * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne12,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src1 = src1 + offset1;
+    dst  = dst  + offsetd;
+
+    int ix = get_sub_group_local_id()/8;  // super block index
+    int it = get_sub_group_local_id()%8;  // sub block inside the super block
+
+    int nb = ne00/QK_K;
+
+    int r0 = get_group_id(0);
+    int r1 = get_group_id(1);
+    int im = get_group_id(2);
+    int first_row = (r0 * N_SIMDGROUP + get_sub_group_id()) * N_DST;
+
+    int i12 = im%ne12;
+    int i13 = im/ne12;
+
+    ulong offset_src0 = (offset0 + first_row*nb01 + (i12/r2)*nb02 + (i13/r3)*nb03)/sizeof(block_iq3_xxs);
+    ulong row_stride = nb01/sizeof(block_iq3_xxs);
+    ulong offset_src1 =        r1*nb11 + (i12   )*nb12 + (i13   )*nb13;
+
+    global float         * y = (global float         *) (src1 + offset_src1);
+
+    float yl[32];
+    float sumf[N_DST] = {0.f};
+    float all_sum;
+
+    global float * y4 = y + ix * QK_K + 32 * it;
+
+    for (int ib = ix; ib < nb; ib += BLOCK_STRIDE) {
+        for (int i = 0; i < 32; ++i) {
+            yl[i] = y4[i];
+        }
+
+        // only ne01 output rows exist; reading past them can pull in an inf scale
+        for (int row = 0; row < N_DST && first_row + row < ne01; row++) {
+            const ulong b = offset_src0 + ib + row*row_stride;
+
+            global uchar * qs = src0_qs + 64*b + 8*it;
+
+            uint aux32 = src0_sas[8*b + it];
+
+            float db = (float)src0_d[b] * (0.5f + (float)(aux32 >> 28)) * 0.5f;
+
+            float acc = 0.f;
+            for (int l = 0; l < 4; ++l) {
+                uchar signs = iq3_xxs_signs((aux32 >> (7*l)) & 127u);
+                uint  g1 = iq3_xxs_grid[qs[2*l+0]];
+                uint  g2 = iq3_xxs_grid[qs[2*l+1]];
+
+                for (int j = 0; j < 4; ++j) {
+                    float s1 = (signs & (1 << j))       ? -1.f : 1.f;
+                    float s2 = (signs & (1 << (j+4)))   ? -1.f : 1.f;
+                    acc += yl[8*l+j+0] * (float)((g1 >> (8*j)) & 0xFF) * s1;
+                    acc += yl[8*l+j+4] * (float)((g2 >> (8*j)) & 0xFF) * s2;
+                }
+            }
+
+            sumf[row] += db * acc;
+        }
+
+        y4 += BLOCK_STRIDE * QK_K;
+    }
+
+    global float * dst_f32 = (global float *) dst + im*ne0*ne1 + r1*ne0;
+
+    for (int row = 0; row < N_DST; ++row) {
+        all_sum = sub_group_reduce_add(sumf[row]);
+        if (first_row + row < ne01) {
+            if (get_sub_group_local_id() == 0) {
+                dst_f32[first_row + row] = all_sum;
+            }
+        }
+    }
+}

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

@@ -0,0 +1,134 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_intel_required_subgroup_size
+#pragma OPENCL EXTENSION cl_intel_required_subgroup_size : enable
+#define INTEL_GPU 1
+#define REQD_SUBGROUP_SIZE_16 __attribute__((intel_reqd_sub_group_size(16)))
+#define REQD_SUBGROUP_SIZE_32 __attribute__((intel_reqd_sub_group_size(32)))
+#elif defined(cl_qcom_reqd_sub_group_size)
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+#define ADRENO_GPU 1
+#define REQD_SUBGROUP_SIZE_64  __attribute__((qcom_reqd_sub_group_size("half")))
+#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+
+#define QK_K 256
+
+typedef struct {
+    half d;
+    ushort scales_h;
+    uchar scales_l[QK_K/64];
+    uchar qs[QK_K/2];
+} block_iq4_xs;
+
+#undef N_DST
+#undef N_SIMDGROUP
+#undef N_SIMDWIDTH
+
+#ifdef INTEL_GPU
+#define N_DST 4
+#define N_SIMDGROUP 1
+#define N_SIMDWIDTH 16
+#elif defined (ADRENO_GPU)
+#define N_DST 4
+#define N_SIMDGROUP 1
+#define N_SIMDWIDTH 64
+#endif
+
+#undef  BLOCK_STRIDE
+// 8 threads cover one super block, one 32 element sub block each
+#define BLOCK_STRIDE (N_SIMDWIDTH/8)
+
+#ifdef INTEL_GPU
+REQD_SUBGROUP_SIZE_16
+#elif defined (ADRENO_GPU)
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mv_iq4_xs_f32(
+        global char * src0,
+        ulong offset0,
+        global char * src1,
+        ulong offset1,
+        global char * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne12,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src0 = src0 + offset0;
+    src1 = src1 + offset1;
+    dst  = dst  + offsetd;
+
+    int ix = get_sub_group_local_id()/8;  // super block index
+    int it = get_sub_group_local_id()%8;  // sub block inside the super block
+
+    int nb = ne00/QK_K;
+
+    int r0 = get_group_id(0);
+    int r1 = get_group_id(1);
+    int im = get_group_id(2);
+    int first_row = (r0 * N_SIMDGROUP + get_sub_group_id()) * N_DST;
+
+    int i12 = im%ne12;
+    int i13 = im/ne12;
+
+    ulong offset_src0 = first_row*nb01 + (i12/r2)*nb02 + (i13/r3)*nb03;
+    ulong offset_src1 =        r1*nb11 + (i12   )*nb12 + (i13   )*nb13;
+
+    global block_iq4_xs * x = (global block_iq4_xs *) (src0 + offset_src0);
+    global float         * y = (global float         *) (src1 + offset_src1);
+
+    float yl[32];
+    float sumf[N_DST] = {0.f};
+    float all_sum;
+
+    global float * y4 = y + ix * QK_K + 32 * it;
+
+    for (int ib = ix; ib < nb; ib += BLOCK_STRIDE) {
+        for (int i = 0; i < 32; ++i) {
+            yl[i] = y4[i];
+        }
+
+        global char * xrow = (global char *)(x + ib);
+
+        // only ne01 output rows exist; reading past them can pull in an inf scale
+        for (int row = 0; row < N_DST && first_row + row < ne01; row++) {
+            global block_iq4_xs * xb = (global block_iq4_xs *)(xrow + row*nb01);
+
+            int ls = ((xb->scales_l[it/2] >> (4*(it%2))) & 15) | (((xb->scales_h >> (2*it)) & 3) << 4);
+            float db = (float)xb->d * (float)(ls - 32);
+            global uchar * qs = xb->qs + 16*it;
+
+            float acc = 0.f;
+            for (int j = 0; j < 16; ++j) {
+                acc += yl[j]    * (float)iq4_xs_values[qs[j] & 15];
+                acc += yl[j+16] * (float)iq4_xs_values[qs[j] >>  4];
+            }
+
+            sumf[row] += db * acc;
+        }
+
+        y4 += BLOCK_STRIDE * QK_K;
+    }
+
+    global float * dst_f32 = (global float *) dst + im*ne0*ne1 + r1*ne0;
+
+    for (int row = 0; row < N_DST; ++row) {
+        all_sum = sub_group_reduce_add(sumf[row]);
+        if (first_row + row < ne01) {
+            if (get_sub_group_local_id() == 0) {
+                dst_f32[first_row + row] = all_sum;
+            }
+        }
+    }
+}

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

@@ -0,0 +1,134 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_intel_required_subgroup_size
+#pragma OPENCL EXTENSION cl_intel_required_subgroup_size : enable
+#define INTEL_GPU 1
+#define REQD_SUBGROUP_SIZE_16 __attribute__((intel_reqd_sub_group_size(16)))
+#define REQD_SUBGROUP_SIZE_32 __attribute__((intel_reqd_sub_group_size(32)))
+#elif defined(cl_qcom_reqd_sub_group_size)
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+#define ADRENO_GPU 1
+#define REQD_SUBGROUP_SIZE_64  __attribute__((qcom_reqd_sub_group_size("half")))
+#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+
+#define QK_K 256
+
+typedef struct {
+    half d;
+    ushort scales_h;
+    uchar scales_l[QK_K/64];
+    uchar qs[QK_K/2];
+} block_iq4_xs;
+
+#undef N_DST
+#undef N_SIMDGROUP
+#undef N_SIMDWIDTH
+
+#ifdef INTEL_GPU
+#define N_DST 4
+#define N_SIMDGROUP 1
+#define N_SIMDWIDTH 16
+#elif defined (ADRENO_GPU)
+#define N_DST 4
+#define N_SIMDGROUP 1
+#define N_SIMDWIDTH 64
+#endif
+
+#undef  BLOCK_STRIDE
+// 8 threads cover one super block, one 32 element sub block each
+#define BLOCK_STRIDE (N_SIMDWIDTH/8)
+
+#ifdef INTEL_GPU
+REQD_SUBGROUP_SIZE_16
+#elif defined (ADRENO_GPU)
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mv_iq4_xs_f32_flat(
+        global uchar * src0_qs,
+        global uchar * src0_sl,
+        global ushort * src0_sh,
+        global half * src0_d,
+        ulong offset0,
+        global char * src1,
+        ulong offset1,
+        global char * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne12,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src1 = src1 + offset1;
+    dst  = dst  + offsetd;
+
+    int ix = get_sub_group_local_id()/8;  // super block index
+    int it = get_sub_group_local_id()%8;  // sub block inside the super block
+
+    int nb = ne00/QK_K;
+
+    int r0 = get_group_id(0);
+    int r1 = get_group_id(1);
+    int im = get_group_id(2);
+    int first_row = (r0 * N_SIMDGROUP + get_sub_group_id()) * N_DST;
+
+    int i12 = im%ne12;
+    int i13 = im/ne12;
+
+    ulong offset_src0 = (offset0 + first_row*nb01 + (i12/r2)*nb02 + (i13/r3)*nb03)/sizeof(block_iq4_xs);
+    ulong row_stride = nb01/sizeof(block_iq4_xs);
+    ulong offset_src1 =        r1*nb11 + (i12   )*nb12 + (i13   )*nb13;
+
+    global float         * y = (global float         *) (src1 + offset_src1);
+
+    float yl[32];
+    float sumf[N_DST] = {0.f};
+    float all_sum;
+
+    global float * y4 = y + ix * QK_K + 32 * it;
+
+    for (int ib = ix; ib < nb; ib += BLOCK_STRIDE) {
+        for (int i = 0; i < 32; ++i) {
+            yl[i] = y4[i];
+        }
+
+        // only ne01 output rows exist; reading past them can pull in an inf scale
+        for (int row = 0; row < N_DST && first_row + row < ne01; row++) {
+            const ulong b = offset_src0 + ib + row*row_stride;
+
+            int ls = ((src0_sl[4*b + it/2] >> (4*(it%2))) & 15) | (((src0_sh[b] >> (2*it)) & 3) << 4);
+            float db = (float)src0_d[b] * (float)(ls - 32);
+            global uchar * qs = src0_qs + 128*b + 16*it;
+
+            float acc = 0.f;
+            for (int j = 0; j < 16; ++j) {
+                acc += yl[j]    * (float)iq4_xs_values[qs[j] & 15];
+                acc += yl[j+16] * (float)iq4_xs_values[qs[j] >>  4];
+            }
+
+            sumf[row] += db * acc;
+        }
+
+        y4 += BLOCK_STRIDE * QK_K;
+    }
+
+    global float * dst_f32 = (global float *) dst + im*ne0*ne1 + r1*ne0;
+
+    for (int row = 0; row < N_DST; ++row) {
+        all_sum = sub_group_reduce_add(sumf[row]);
+        if (first_row + row < ne01) {
+            if (get_sub_group_local_id() == 0) {
+                dst_f32[first_row + row] = all_sum;
+            }
+        }
+    }
+}

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

@@ -0,0 +1,129 @@
+// Integer storage preserves the mask bits without FP16 conversion.
+#if MASK_F16
+#define MASK_T ushort
+#define MASK_V ushort8
+#define MASK_NV 8
+#define MASK_NINF 0xfc00
+#define MASK_STORE vstore8
+#else
+#define MASK_T uint
+#define MASK_V uint4
+#define MASK_NV 4
+#define MASK_NINF 0xff800000u
+#define MASK_STORE vstore4
+#endif
+
+kernel void kernel_top_k_mask_fill(
+        global char * dst,
+        ulong         offsetd,
+        ulong         n
+) {
+    global MASK_T * out = (global MASK_T *) (dst + offsetd);
+    const ulong c = get_global_id(0)*MASK_NV;
+    if (c + MASK_NV <= n) {
+        MASK_STORE((MASK_V) (MASK_NINF), 0, out + c);
+    } else {
+        for (ulong j = c; j < n; ++j) {
+            out[j] = MASK_NINF;
+        }
+    }
+}
+
+kernel void kernel_top_k_mask(
+        global const char * mask,
+        ulong               offsetm,
+        global const char * top_k,
+        ulong               offsetk,
+        global const char * score,
+        ulong               offsets,
+        global const char * blocks,
+        ulong               offsetb,
+        global const char * extra,
+        ulong               offsete,
+        global       char * dst,
+        ulong               offsetd,
+        ulong nem0,
+        ulong nem1,
+        ulong nem2,
+        ulong nem3,
+        ulong nbm1,
+        ulong nbm2,
+        ulong nbm3,
+        ulong nek0,
+        ulong nbk0,
+        ulong nbk1,
+        ulong nbk2,
+        ulong nbk3,
+        ulong nes0,
+        ulong nbs0,
+        ulong nbs1,
+        ulong nbs2,
+        ulong nbs3,
+        ulong neb0,
+        ulong neb1,
+        ulong nbb0,
+        ulong nbb1,
+        ulong nee0,
+        ulong nbe0,
+        ulong nbe1,
+        ulong nbe2,
+        ulong nbe3,
+        ulong ned1,
+        ulong ned2,
+        ulong nbd1,
+        ulong nbd2,
+        ulong nbd3,
+        int   fill
+) {
+    const ulong row = get_group_id(1);
+    const ulong i1 = row % ned1;
+    const ulong i2 = (row / ned1) % ned2;
+    const ulong i3 = row / (ned1*ned2);
+    const ulong tid = get_local_id(0);
+    const ulong nth = get_local_size(0);
+    const ulong n_pick = blocks ? neb1 : score ? nes0 : nem0;
+
+    global const MASK_T * mask_row = (global const MASK_T *) (mask + offsetm
+        + (i1 % nem1)*nbm1 + (i2 % nem2)*nbm2 + (i3 % nem3)*nbm3);
+    global MASK_T * dst_row = (global MASK_T *) (dst + offsetd + i1*nbd1 + i2*nbd2 + i3*nbd3);
+    global const char * top_k_row = top_k + offsetk + i1*nbk1 + i2*nbk2 + i3*nbk3;
+
+    if (fill) {
+        for (ulong c = tid; c < nem0/MASK_NV; c += nth) {
+            MASK_STORE((MASK_V) (MASK_NINF), c, dst_row);
+        }
+        for (ulong c = nem0/MASK_NV*MASK_NV + tid; c < nem0; c += nth) {
+            dst_row[c] = MASK_NINF;
+        }
+        barrier(CLK_GLOBAL_MEM_FENCE);
+    }
+
+    // Repeated cells copy the same mask value, as in the CUDA and Metal kernels.
+    for (ulong j = tid; j < nek0 + nee0; j += nth) {
+        if (j >= nek0) {
+            const int c = *(global const int *) (extra + offsete + (j - nek0)*nbe0 + i1*nbe1 + i2*nbe2 + i3*nbe3);
+            if (c >= 0 && (ulong) c < nem0) {
+                dst_row[c] = mask_row[c];
+            }
+            continue;
+        }
+
+        const int p = *(global const int *) (top_k_row + j*nbk0);
+        if (p < 0 || (ulong) p >= n_pick) {
+            continue;
+        }
+        if (score && !(*(global const float *) (score + offsets + p*nbs0 + i1*nbs1 + i2*nbs2 + i3*nbs3) > -INFINITY)) {
+            continue;
+        }
+        if (blocks) {
+            for (ulong k = 0; k < neb0; ++k) {
+                const int c = *(global const int *) (blocks + offsetb + k*nbb0 + p*nbb1);
+                if (c >= 0 && (ulong) c < nem0) {
+                    dst_row[c] = mask_row[c];
+                }
+            }
+        } else if ((ulong) p < nem0) {
+            dst_row[p] = mask_row[p];
+        }
+    }
+}

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

@@ -1520,6 +1520,11 @@ static ggml_openvino_op_support is_op_supported_case(const ggml_tensor * op) {
             return {false, "FLASH_ATTN_EXT gemma3n pattern on GPU is not supported"};
         }
 
+        // K/V row lists (ggml_flash_attn_ext_rows) are not supported
+        if (op->src[5] != nullptr) {
+            return {false, "FLASH_ATTN_EXT with K/V row lists is not supported"};
+        }
+
         if (op->src[4] != nullptr) {
             return {false, "FLASH_ATTN_EXT with sinks is not supported"};
         }

ggml/src/ggml-rpc/ggml-rpc.cpp

@@ -1430,6 +1430,9 @@ public:
 
 private:
     void sync_all_backends();
+    void comm_free_device(uint32_t dev_id);
+    static bool comm_allreduce_pairwise(void * comm_ctx_v, ggml_tensor ** tensors);
+    static bool comm_bootstrap_exchange(void * user_data, const void * send_buf, void * recv_buf, size_t size);
     bool get_cached_file(uint64_t hash, std::vector<uint8_t> & data);
     ggml_tensor * deserialize_tensor(struct ggml_context * ctx, const rpc_tensor * tensor);
     ggml_tensor * create_node(uint64_t id,
@@ -1438,15 +1441,26 @@ private:
                               std::unordered_map<uint64_t, struct ggml_tensor*> & tensor_map);
 
 
-    // pairwise allreduce over a direct connection to the peer server
+    // Per-device state for one communicator.
     struct comm_state {
-        socket_ptr              peer;
+        // Sockets to the other servers. Rank 0 holds one per peer in rank order, so peers[i] is
+        // rank i + 1. Every other rank holds one to rank 0. The pairwise path reduces over peers[0].
+        std::vector<socket_ptr> peers;
         uint32_t                rank = 0;
         uint32_t                world = 0;
+        ggml_backend_t          pairwise_backend = nullptr;
         ggml_backend_buffer_ptr scratch;
         size_t                  scratch_size = 0;
         std::vector<uint8_t>    send_buf;
         std::vector<uint8_t>    recv_buf;
+
+        // The backend's communicator. Null on the pairwise path, which has none.
+        void *                              backend_comm_ctx = nullptr;
+
+        // The chosen reduction, set once a path is agreed. Receives backend_comm_ctx, or
+        // comm_state on the pairwise path. A communicator covers one device, so the caller
+        // passes a one-element tensor array.
+        ggml_backend_comm_allreduce_tensor_t allreduce_fn    = nullptr;
     };
 
     std::vector<ggml_backend_t> backends;
@@ -2225,78 +2239,188 @@ void rpc_server::sync_all_backends() {
     }
 }
 
-// The comm link between two servers uses the same caps negotiation as the client HELLO,
+void rpc_server::comm_free_device(uint32_t dev_id) {
+    comm_state & state = comm_states[dev_id];
+    if (state.backend_comm_ctx != nullptr) {
+        auto comm_free = (ggml_backend_comm_free_t) ggml_backend_reg_get_proc_address(
+            ggml_backend_dev_backend_reg(ggml_backend_get_device(backends[dev_id])), "ggml_backend_comm_free");
+        if (comm_free != nullptr) {
+            comm_free(state.backend_comm_ctx);
+        }
+    }
+    state = comm_state();
+}
+
+// Exchange RPC transport capabilities over a peer connection. recv_first selects which side
+// receives first. Both sides receiving first would deadlock.
+static bool comm_caps_handshake(const socket_ptr & peer, bool recv_first) {
+    uint8_t local_caps[RPC_CONN_CAPS_SIZE]  = {};
+    uint8_t remote_caps[RPC_CONN_CAPS_SIZE] = {};
+    peer->get_caps(local_caps);
+    if (recv_first) {
+        if (!peer->recv_data(remote_caps, sizeof(remote_caps)) ||
+            !peer->send_data(local_caps, sizeof(local_caps))) {
+            return false;
+        }
+    } else {
+        if (!peer->send_data(local_caps, sizeof(local_caps)) ||
+            !peer->recv_data(remote_caps, sizeof(remote_caps))) {
+            return false;
+        }
+    }
+    peer->update_caps(remote_caps);
+    return true;
+}
+
+// The comm link between servers uses the same caps negotiation as the client HELLO,
 // so it gets the same transport upgrades (e.g. RDMA).
 bool rpc_server::comm_init(const rpc_msg_comm_init_req & request, rpc_msg_comm_init_rsp & response) {
     response.ok = 0;
-    if (request.device >= backends.size() || request.world != 2 || request.rank >= request.world) {
+    if (request.device >= backends.size() || request.world < 2 || request.rank >= request.world) {
         return true;
     }
-    comm_state & state = comm_states[request.device];
-    if (state.peer != nullptr) {
+    const uint32_t dev_id = request.device;
+    comm_state & state = comm_states[dev_id];
+    if (state.allreduce_fn != nullptr) {
         response.ok = 1;
         return true;
     }
-    uint8_t local_caps[RPC_CONN_CAPS_SIZE] = {};
-    uint8_t remote_caps[RPC_CONN_CAPS_SIZE] = {};
+
+    // Joining a communicator blocks until every rank arrives, so drain outstanding work first.
+    sync_all_backends();
+
+    // Rank 0 listens and accepts world-1 peers; every other rank connects back to it.
     if (request.rank == 0) {
         socket_ptr srv = socket_t::create_server("0.0.0.0", request.port);
         if (srv == nullptr) {
             LOG_ERROR("[%s] failed to listen on comm port %u\n", __func__, request.port);
             return true;
         }
-        state.peer = srv->accept();
-        if (state.peer == nullptr) {
-            return true;
+        // Store peers[i] as rank i + 1, as comm_bootstrap_exchange needs.
+        // Each peer sends its rank because peers may connect in any order.
+        state.peers.resize(request.world - 1);
+        for (uint32_t i = 1; i < request.world; i++) {
+            socket_ptr peer = srv->accept();
+            uint32_t peer_rank = 0;
+            if (peer == nullptr || !comm_caps_handshake(peer, /*recv_first =*/ true) ||
+                !peer->recv_data(&peer_rank, sizeof(peer_rank)) ||
+                peer_rank == 0 || peer_rank >= request.world || state.peers[peer_rank - 1] != nullptr) {
+                LOG_ERROR("[%s] bad peer connection (rank %u)\n", __func__, peer_rank);
+                state.peers.clear();
+                return true;
+            }
+            state.peers[peer_rank - 1] = std::move(peer);
         }
-        if (!state.peer->recv_data(remote_caps, sizeof(remote_caps))) {
-            state.peer = nullptr;
-            return true;
-        }
-        state.peer->get_caps(local_caps);
-        if (!state.peer->send_data(local_caps, sizeof(local_caps))) {
-            state.peer = nullptr;
-            return true;
-        }
-        state.peer->update_caps(remote_caps);
     } else {
         const std::string host(request.host, strnlen(request.host, sizeof(request.host)));
+        socket_ptr peer;
         // rank 0 may not be listening yet, retry for a few seconds
-        for (int i = 0; i < 100 && state.peer == nullptr; i++) {
-            state.peer = socket_t::connect(host.c_str(), request.port);
-            if (state.peer == nullptr) {
+        for (int i = 0; i < 100 && peer == nullptr; i++) {
+            peer = socket_t::connect(host.c_str(), request.port);
+            if (peer == nullptr) {
                 std::this_thread::sleep_for(std::chrono::milliseconds(50));
             }
         }
-        if (state.peer == nullptr) {
+        if (peer == nullptr) {
             LOG_ERROR("[%s] failed to connect to peer %s:%u\n", __func__, host.c_str(), request.port);
             return true;
         }
-        state.peer->get_caps(local_caps);
-        if (!state.peer->send_data(local_caps, sizeof(local_caps)) ||
-            !state.peer->recv_data(remote_caps, sizeof(remote_caps))) {
-            state.peer = nullptr;
+        if (!comm_caps_handshake(peer, /*recv_first =*/ false) ||
+            !peer->send_data(&request.rank, sizeof(request.rank)) || !peer->flush()) {
             return true;
         }
-        state.peer->update_caps(remote_caps);
+        state.peers.push_back(std::move(peer));
     }
+
     state.rank  = request.rank;
     state.world = request.world;
-    LOG_INFO("[%s] device %u joined pairwise comm as rank %u\n", __func__, request.device, request.rank);
+
+    // Forces the pairwise path, for comparing the two on one build. Setting it on any one rank
+    // is enough, because the negotiation spreads the decision.
+    const bool no_backend_communicator = std::getenv("GGML_RPC_NO_BACKEND_COMMUNICATOR") != nullptr;
+
+    ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backends[dev_id]));
+    auto init_rank    = (ggml_backend_comm_init_rank_t)
+        ggml_backend_reg_get_proc_address(reg, "ggml_backend_comm_init_rank");
+    auto allreduce_fn = (ggml_backend_comm_allreduce_tensor_t)
+        ggml_backend_reg_get_proc_address(reg, "ggml_backend_comm_allreduce_tensor");
+
+    // The ranks must agree on backend communicator or pairwise. Every rank shares whether its
+    // backend provides a communicator, and all ranks use it only if every rank can.
+    const uint8_t self_can_use_backend_comm = !no_backend_communicator && init_rank != nullptr && allreduce_fn != nullptr;
+    std::vector<uint8_t> all_can_use_backend_comm(request.world);
+    if (!comm_bootstrap_exchange(&state, &self_can_use_backend_comm, all_can_use_backend_comm.data(),
+                                 sizeof(self_can_use_backend_comm))) {
+        state.peers.clear();
+        return true;
+    }
+    const bool use_backend_communicator = std::all_of(all_can_use_backend_comm.begin(), all_can_use_backend_comm.end(),
+                                                      [](uint8_t c) { return c != 0; });
+
+    if (use_backend_communicator) {
+        void * backend_ctx = init_rank(backends[dev_id], (int) request.rank, (int) request.world, comm_bootstrap_exchange, &state);
+        if (backend_ctx == nullptr) {
+            LOG_ERROR("[%s] failed to join communicator as rank %u/%u\n", __func__, request.rank, request.world);
+            state.peers.clear();
+            return true;
+        }
+        state.backend_comm_ctx = backend_ctx;
+        state.allreduce_fn     = allreduce_fn;
+        // The sockets have done their job; the collective runs between the backends from here on.
+        state.peers.clear();
+        LOG_INFO("[%s] device %u using backend communicator, rank %u/%u\n", __func__, dev_id, request.rank, request.world);
+        response.ok = 1;
+        return true;
+    }
+
+    // No backend communicator, so fall back to the pairwise exchange, which only covers two ranks.
+    if (request.world != 2) {
+        LOG_ERROR("[%s] world size %u requires a backend communicator\n", __func__, request.world);
+        state.peers.clear();
+        return true;
+    }
+    state.pairwise_backend = backends[dev_id];
+    state.allreduce_fn     = comm_allreduce_pairwise;
+    LOG_INFO("[%s] device %u using pairwise communicator, rank %u\n", __func__, dev_id, request.rank);
     response.ok = 1;
     return true;
 }
 
-bool rpc_server::comm_allreduce(const rpc_msg_comm_allreduce_req & request) {
-    if (request.device >= backends.size()) {
-        return false;
+// Bootstrap callback for ggml_backend_comm_init_rank_t, over the peer sockets. Rank 0 collects
+// size bytes from every rank in rank order, then sends the whole array back to every rank.
+// user_data is the comm_state.
+bool rpc_server::comm_bootstrap_exchange(void * user_data, const void * send_buf, void * recv_buf, size_t size) {
+    comm_state & state = *(comm_state *) user_data;
+    uint8_t * recv = (uint8_t *) recv_buf;
+    const size_t total = (size_t) state.world * size;
+
+    if (state.rank == 0) {
+        memcpy(recv, send_buf, size);
+        for (size_t i = 0; i < state.peers.size(); i++) {
+            if (!state.peers[i]->recv_data(recv + (i + 1) * size, size)) {
+                return false;
+            }
+        }
+        for (const socket_ptr & peer : state.peers) {
+            if (!peer->send_data(recv, total) || !peer->flush()) {
+                return false;
+            }
+        }
+        return true;
     }
-    comm_state & state = comm_states[request.device];
-    if (state.peer == nullptr) {
-        LOG_ERROR("[%s] no communicator for device %u\n", __func__, request.device);
-        return false;
-    }
-    ggml_backend_t backend = backends[request.device];
+    return state.peers[0]->send_data(send_buf, size) && state.peers[0]->flush() &&
+           state.peers[0]->recv_data(recv, total);
+}
+
+// Reduces this rank's contribution in place, exchanging partials with the peer over peers[0].
+// Used when the backend provides no communicator. comm_ctx_v is the comm_state.
+bool rpc_server::comm_allreduce_pairwise(void * comm_ctx_v, ggml_tensor ** tensors) {
+    comm_state & state = *(comm_state *) comm_ctx_v;
+    ggml_tensor * t_dst = tensors[0];
+    ggml_backend_t backend = state.pairwise_backend;
+
+    const size_t  nbytes = ggml_nbytes(t_dst);
+    const int64_t ne     = ggml_nelements(t_dst);
 
     size_t ctx_size = 16*ggml_tensor_overhead() + 2*ggml_graph_overhead_custom(8, false);
     struct ggml_init_params params = {
@@ -2307,16 +2431,7 @@ bool rpc_server::comm_allreduce(const rpc_msg_comm_allreduce_req & request) {
     ggml_context_ptr ctx_ptr { ggml_init(params) };
     GGML_ASSERT(ctx_ptr != nullptr);
     ggml_context * ctx = ctx_ptr.get();
-    ggml_tensor * t_dst = deserialize_tensor(ctx, &request.tensor);
-    if (t_dst == nullptr || t_dst->buffer == nullptr) {
-        LOG_ERROR("[%s] error deserializing tensor\n", __func__);
-        return false;
-    }
-    const size_t  nbytes = ggml_nbytes(t_dst);
-    const int64_t ne     = ggml_nelements(t_dst);
-    if (nbytes == 0) {
-        return true;
-    }
+
     // reduce large partials in bf16 to halve the wire bytes; small (decode-sized) ones
     // stay f32 since the extra casts and sync cost more than the bytes saved
     const bool   wire_bf16  = t_dst->type == GGML_TYPE_F32 && ne >= 32768;
@@ -2372,13 +2487,13 @@ bool rpc_server::comm_allreduce(const rpc_msg_comm_allreduce_req & request) {
 
     // rank 0 sends first, rank 1 receives first, so large payloads cannot deadlock
     if (state.rank == 0) {
-        if (!state.peer->send_data(state.send_buf.data(), wire_bytes) || !state.peer->flush() ||
-            !state.peer->recv_data(state.recv_buf.data(), wire_bytes)) {
+        if (!state.peers[0]->send_data(state.send_buf.data(), wire_bytes) || !state.peers[0]->flush() ||
+            !state.peers[0]->recv_data(state.recv_buf.data(), wire_bytes)) {
             return false;
         }
     } else {
-        if (!state.peer->recv_data(state.recv_buf.data(), wire_bytes) ||
-            !state.peer->send_data(state.send_buf.data(), wire_bytes) || !state.peer->flush()) {
+        if (!state.peers[0]->recv_data(state.recv_buf.data(), wire_bytes) ||
+            !state.peers[0]->send_data(state.send_buf.data(), wire_bytes) || !state.peers[0]->flush()) {
             return false;
         }
     }
@@ -2408,11 +2523,50 @@ bool rpc_server::comm_allreduce(const rpc_msg_comm_allreduce_req & request) {
     return true;
 }
 
+bool rpc_server::comm_allreduce(const rpc_msg_comm_allreduce_req & request) {
+    if (request.device >= backends.size()) {
+        return false;
+    }
+    comm_state & state = comm_states[request.device];
+    if (state.allreduce_fn == nullptr) {
+        LOG_ERROR("[%s] no communicator for device %u\n", __func__, request.device);
+        return false;
+    }
+
+    struct ggml_init_params params = {
+        /*.mem_size   =*/ ggml_tensor_overhead(),
+        /*.mem_buffer =*/ NULL,
+        /*.no_alloc   =*/ true,
+    };
+    ggml_context_ptr ctx_ptr { ggml_init(params) };
+    GGML_ASSERT(ctx_ptr != nullptr);
+    ggml_tensor * t_dst = deserialize_tensor(ctx_ptr.get(), &request.tensor);
+    if (t_dst == nullptr || t_dst->buffer == nullptr) {
+        LOG_ERROR("[%s] error deserializing tensor\n", __func__);
+        return false;
+    }
+    if (ggml_nbytes(t_dst) == 0) {
+        return true;
+    }
+
+    // Neither path needs a wait here. The pairwise path waits internally, and the backend orders
+    // its reduction after the work that produced the partial sum.
+    void * comm_ctx = state.backend_comm_ctx != nullptr ? state.backend_comm_ctx : &state;
+    if (!state.allreduce_fn(comm_ctx, &t_dst)) {
+        LOG_ERROR("[%s] allreduce failed for tensor %s\n", __func__, t_dst->name);
+        return false;
+    }
+    return true;
+}
+
 bool rpc_server::comm_free(const rpc_msg_comm_free_req & request) {
     if (request.device >= backends.size()) {
         return false;
     }
-    comm_states[request.device] = comm_state();
+    // Wait for outstanding work. Destroying a communicator while its kernels are still running
+    // is unsafe.
+    sync_all_backends();
+    comm_free_device(request.device);
     return true;
 }
 
@@ -2432,6 +2586,14 @@ bool rpc_server::get_device_memory(const rpc_msg_get_device_memory_req & request
 }
 
 rpc_server::~rpc_server() {
+    // Wait for outstanding work. Destroying a communicator while its kernels are still running
+    // is unsafe.
+    sync_all_backends();
+
+    // Free any communicator still open. A client can disconnect without sending COMM_FREE.
+    for (uint32_t dev_id = 0; dev_id < comm_states.size(); dev_id++) {
+        comm_free_device(dev_id);
+    }
     for (auto buffer : buffers) {
         ggml_backend_buffer_free(buffer);
     }
@@ -2997,7 +3159,7 @@ static void ggml_backend_rpc_comm_free(void * comm_ctx_v) {
 }
 
 static void * ggml_backend_rpc_comm_init(ggml_backend_t * backends, size_t n_backends) {
-    if (n_backends != 2 || std::getenv("GGML_RPC_NO_COMM") != nullptr) {
+    if (n_backends < 2 || std::getenv("GGML_RPC_NO_COMM") != nullptr) {
         return nullptr;
     }
     std::vector<ggml_backend_rpc_comm_context::rank_info> ranks;

ggml/src/ggml-sycl/CMakeLists.txt

@@ -107,24 +107,36 @@ target_compile_options(ggml-sycl PRIVATE "-Wno-narrowing")
 
 message(STATUS "GGML_SYCL_SUPPORT_LEVEL_ZERO_API ${GGML_SYCL_SUPPORT_LEVEL_ZERO_API}")
 if (GGML_SYCL_SUPPORT_LEVEL_ZERO_API)
-    # Link against Level Zero loader for direct device memory allocation.
-    # Avoids sycl::malloc_device triggering DMA-buf/TTM system RAM staging
-    # in the xe kernel driver during multi-GPU inference.
-    find_path(LEVEL_ZERO_DEV_INCLUDE_DIR level_zero/ze_api.h HINTS ${ONEAPI_ROOT}/include ${LEVEL_ZERO_V1_SDK_PATH}/include)
-    find_library(ZE_LOADER_LIB ze_loader HINTS ${ONEAPI_ROOT}/lib ${LEVEL_ZERO_V1_SDK_LIB_PATH} ENV LD_LIBRARY_PATH)
-    if(ZE_LOADER_LIB AND LEVEL_ZERO_DEV_INCLUDE_DIR)
-        target_link_libraries(ggml-sycl PRIVATE ${ZE_LOADER_LIB})
-        target_compile_definitions(ggml-sycl PRIVATE GGML_SYCL_SUPPORT_LEVEL_ZERO_API)
-        message(STATUS "Level Zero loader found: ${ZE_LOADER_LIB}")
-        message(STATUS "Level Zero development headers found: ${LEVEL_ZERO_DEV_INCLUDE_DIR}")
+    find_package(PkgConfig)
+    if (PkgConfig_FOUND)
+        pkg_check_modules(LEVEL_ZERO level-zero)
+        if (LEVEL_ZERO_FOUND)
+            target_include_directories(ggml-sycl PRIVATE ${LEVEL_ZERO_INCLUDE_DIRS})
+            target_link_libraries(ggml-sycl PRIVATE ${LEVEL_ZERO_LINK_LIBRARIES})
+            target_compile_definitions(ggml-sycl PRIVATE GGML_SYCL_SUPPORT_LEVEL_ZERO_API)
+            message(STATUS "Level Zero loader found: ${LEVEL_ZERO_LINK_LIBRARIES}")
+            message(STATUS "Level Zero development headers found: ${LEVEL_ZERO_INCLUDE_DIRS}")
+        endif()
     else()
-        message(WARNING "Level Zero loader or development headers not found, "
-                        "Level Zero API support disabled. "
-                        "Please install the Level Zero SDK/development package "
-                        "to support Level Zero API features. "
-                        "Level Zero API is not mandatory for SYCL backend, "
-                        "but it is required by the special features for better "
-                        "function & performance on Intel GPUs.")
+        # Link against Level Zero loader for direct device memory allocation.
+        # Avoids sycl::malloc_device triggering DMA-buf/TTM system RAM staging
+        # in the xe kernel driver during multi-GPU inference.
+        find_path(LEVEL_ZERO_DEV_INCLUDE_DIR level_zero/ze_api.h HINTS ${ONEAPI_ROOT}/include ${LEVEL_ZERO_V1_SDK_PATH}/include)
+        find_library(ZE_LOADER_LIB ze_loader HINTS ${ONEAPI_ROOT}/lib ${LEVEL_ZERO_V1_SDK_LIB_PATH} ENV LD_LIBRARY_PATH)
+        if(ZE_LOADER_LIB AND LEVEL_ZERO_DEV_INCLUDE_DIR)
+            target_link_libraries(ggml-sycl PRIVATE ${ZE_LOADER_LIB})
+            target_compile_definitions(ggml-sycl PRIVATE GGML_SYCL_SUPPORT_LEVEL_ZERO_API)
+            message(STATUS "Level Zero loader found: ${ZE_LOADER_LIB}")
+            message(STATUS "Level Zero development headers found: ${LEVEL_ZERO_DEV_INCLUDE_DIR}")
+        else()
+            message(WARNING "Level Zero loader or development headers not found, "
+                            "Level Zero API support disabled. "
+                            "Please install the Level Zero SDK/development package "
+                            "to support Level Zero API features. "
+                            "Level Zero API is not mandatory for SYCL backend, "
+                            "but it is required by the special features for better "
+                            "function & performance on Intel GPUs.")
+        endif()
     endif()
 endif()
 

ggml/src/ggml-sycl/dmmv.cpp

@@ -112,11 +112,12 @@ static void dequantize_mul_mat_vec(const void * __restrict__ vx, const dfloat *
     }
 }
 
-template <int qk, int qr, dequantize_kernel_t_reorder dequantize_kernel_reorder>
+template <int qk, int qr, int qbytes, dequantize_kernel_t_reorder dequantize_kernel_reorder>
 static void dequantize_mul_mat_vec_reorder(const void * __restrict__ vx, const dfloat * __restrict__ y, float * __restrict__ dst, const int ncols, const int nrows,
                                    const sycl::nd_item<3> &item_ct1) {
     // qk = quantized weights per x block
     // qr = number of quantized weights per data value in x block
+    // qbytes = bytes of the quants region per block (not qk/qr for bit-packed types)
     const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
                     item_ct1.get_local_id(1);
 
@@ -127,7 +128,7 @@ static void dequantize_mul_mat_vec_reorder(const void * __restrict__ vx, const d
     const int tid = item_ct1.get_local_id(2);
 
 
-    const int ncols_left = ncols % (QK4_0*WARP_SIZE);
+    const int ncols_left = ncols % (qk*WARP_SIZE);
     const int ncols_align = ncols - ncols_left;
     const int iter_stride = 8*2*GGML_SYCL_DMMV_X;
     const int vals_per_iter = iter_stride / WARP_SIZE; // num quantized vals per thread and i iter //64/16=4, 512/16/2= 16
@@ -139,7 +140,7 @@ static void dequantize_mul_mat_vec_reorder(const void * __restrict__ vx, const d
 #else
     float tmp = 0.0f;
 #endif // GGML_SYCL_F16
-    const char *d_ptr = (const char*)vx+ncols*nrows/2;
+    const char *d_ptr = (const char*)vx + (size_t) ncols * nrows / qk * qbytes;
     int i=0;
     for (i = 0; i < ncols_align; i += iter_stride) {
         const int col = i + vals_per_iter*tid;
@@ -155,7 +156,7 @@ static void dequantize_mul_mat_vec_reorder(const void * __restrict__ vx, const d
             // dequantize
             // for qr = 2 the iqs needs to increase by 1 per j iter because 2 weights per data val
             dfloat2 v;
-            dequantize_kernel_reorder((const void *)d_ptr, ib, (const void *)vx, ib * QK4_0 / 2 +iqs+j/qr, v);
+            dequantize_kernel_reorder((const void *)d_ptr, ib, (const void *)vx, ib * (qk/qr) +iqs+j/qr, v);
 
             // matrix multiplication
             // for qr = 2 the y index needs to increase by 1 per j iter because of y_offset = qk/2
@@ -172,7 +173,7 @@ static void dequantize_mul_mat_vec_reorder(const void * __restrict__ vx, const d
     }
 
     for (; i < ncols; i += iter_stride) {
-        if (tid>=ncols_left/QK4_0) continue;
+        if (tid>=ncols_left/qk) continue;
         const int col = i + vals_per_iter*tid;
         const int ib = (row*ncols + col)/qk; // x block index
         const int iqs = (col%qk)/qr; // x quant index
@@ -186,7 +187,7 @@ static void dequantize_mul_mat_vec_reorder(const void * __restrict__ vx, const d
             // dequantize
             // for qr = 2 the iqs needs to increase by 1 per j iter because 2 weights per data val
             dfloat2 v;
-            dequantize_kernel_reorder((const void *)d_ptr, ib, (const void *)vx, ib * QK4_0 / 2 +iqs+j/qr, v);
+            dequantize_kernel_reorder((const void *)d_ptr, ib, (const void *)vx, ib * (qk/qr) +iqs+j/qr, v);
 
             // matrix multiplication
             // for qr = 2 the y index needs to increase by 1 per j iter because of y_offset = qk/2
@@ -1487,10 +1488,10 @@ static void dequantize_mul_mat_vec_q6_k_reorder(const void * __restrict__ vx, co
     }
 }
 
-static void dequantize_mul_mat_vec_q4_0_sycl_reorder(const void *vx, const dfloat *y,
-                                             float *dst, const int ncols,
-                                             const int nrows,
-                                             dpct::queue_ptr stream) {
+template <int qk, int qr, dequantize_kernel_t dequantize_kernel>
+static void launch_dequantize_mul_mat_vec(const void *vx, const dfloat *y, float *dst,
+                                          const int ncols, const int nrows,
+                                          dpct::queue_ptr stream) {
     GGML_ASSERT(ncols % GGML_SYCL_DMMV_X == 0);
     const int block_num_y = (nrows + GGML_SYCL_MMV_Y - 1) / GGML_SYCL_MMV_Y;
     // the number of rows may exceed maximum grid size in the y or z dimensions, use the x dimension instead
@@ -1503,17 +1504,16 @@ static void dequantize_mul_mat_vec_q4_0_sycl_reorder(const void *vx, const dfloa
         stream->parallel_for(
             sycl::nd_range<3>(block_nums * block_dims, block_dims),
             [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                dequantize_mul_mat_vec_reorder<QK4_0, QR4_0, dequantize_q4_0_reorder>(
+                dequantize_mul_mat_vec<qk, qr, dequantize_kernel>(
                     vx, y, dst, ncols, nrows, item_ct1);
             });
     }
 }
 
-
-static void dequantize_mul_mat_vec_q4_0_sycl(const void *vx, const dfloat *y,
-                                             float *dst, const int ncols,
-                                             const int nrows,
-                                             dpct::queue_ptr stream) {
+template <int qk, int qr, int qbytes, dequantize_kernel_t_reorder dequantize_kernel>
+static void launch_dequantize_mul_mat_vec_reorder(const void *vx, const dfloat *y, float *dst,
+                                                  const int ncols, const int nrows,
+                                                  dpct::queue_ptr stream) {
     GGML_ASSERT(ncols % GGML_SYCL_DMMV_X == 0);
     const int block_num_y = (nrows + GGML_SYCL_MMV_Y - 1) / GGML_SYCL_MMV_Y;
     // the number of rows may exceed maximum grid size in the y or z dimensions, use the x dimension instead
@@ -1526,346 +1526,24 @@ static void dequantize_mul_mat_vec_q4_0_sycl(const void *vx, const dfloat *y,
         stream->parallel_for(
             sycl::nd_range<3>(block_nums * block_dims, block_dims),
             [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                dequantize_mul_mat_vec<QK4_0, QR4_0, dequantize_q4_0>(
+                dequantize_mul_mat_vec_reorder<qk, qr, qbytes, dequantize_kernel>(
                     vx, y, dst, ncols, nrows, item_ct1);
             });
     }
 }
 
-static void dequantize_mul_mat_vec_q1_0_sycl_reorder(const void *vx, const dfloat *y,
-                                             float *dst, const int ncols,
-                                             const int nrows,
-                                             dpct::queue_ptr stream) {
-    GGML_ASSERT(ncols % GGML_SYCL_DMMV_X == 0);
-    const int block_num_y = (nrows + GGML_SYCL_MMV_Y - 1) / GGML_SYCL_MMV_Y;
-    // the number of rows may exceed maximum grid size in the y or z dimensions, use the x dimension instead
-    const sycl::range<3> block_nums(1, 1, block_num_y);
-    const sycl::range<3> block_dims(1, GGML_SYCL_MMV_Y, WARP_SIZE);
-    {
-        dpct::has_capability_or_fail(stream->get_device(),
-                                     {sycl::aspect::fp16});
-
-        stream->parallel_for(
-            sycl::nd_range<3>(block_nums * block_dims, block_dims),
-            [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                dequantize_mul_mat_vec_reorder<QK1_0, QR1_0, dequantize_q1_0_reorder>(
-                    vx, y, dst, ncols, nrows, item_ct1);
-            });
-    }
-}
-
-static void dequantize_mul_mat_vec_q1_0_sycl(const void *vx, const dfloat *y,
-                                             float *dst, const int ncols,
-                                             const int nrows,
-                                             dpct::queue_ptr stream) {
-    GGML_ASSERT(ncols % GGML_SYCL_DMMV_X == 0);
-    const int block_num_y = (nrows + GGML_SYCL_MMV_Y - 1) / GGML_SYCL_MMV_Y;
-    // the number of rows may exceed maximum grid size in the y or z dimensions, use the x dimension instead
-    const sycl::range<3> block_nums(1, 1, block_num_y);
-    const sycl::range<3> block_dims(1, GGML_SYCL_MMV_Y, WARP_SIZE);
-    {
-        dpct::has_capability_or_fail(stream->get_device(),
-                                     {sycl::aspect::fp16});
-
-        stream->parallel_for(
-            sycl::nd_range<3>(block_nums * block_dims, block_dims),
-            [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                dequantize_mul_mat_vec<QK1_0, QR1_0, dequantize_q1_0>(
-                    vx, y, dst, ncols, nrows, item_ct1);
-            });
-    }
-}
-
-static void dequantize_mul_mat_vec_q4_1_sycl(const void *vx, const dfloat *y,
-                                             float *dst, const int ncols,
-                                             const int nrows,
-                                             dpct::queue_ptr stream) {
-    GGML_ASSERT(ncols % GGML_SYCL_DMMV_X == 0);
-    const int block_num_y = (nrows + GGML_SYCL_MMV_Y - 1) / GGML_SYCL_MMV_Y;
-    const sycl::range<3> block_nums(1, 1, block_num_y);
-    const sycl::range<3> block_dims(1, GGML_SYCL_MMV_Y, WARP_SIZE);
-    {
-        dpct::has_capability_or_fail(stream->get_device(),
-                                     {sycl::aspect::fp16});
-
-        stream->parallel_for(
-            sycl::nd_range<3>(block_nums * block_dims, block_dims),
-            [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                dequantize_mul_mat_vec<QK4_1, QR4_1, dequantize_q4_1>(
-                    vx, y, dst, ncols, nrows, item_ct1);
-            });
-    }
-}
-
-static void dequantize_mul_mat_vec_q5_0_sycl(const void *vx, const dfloat *y,
-                                             float *dst, const int ncols,
-                                             const int nrows,
-                                             dpct::queue_ptr stream) {
-    GGML_ASSERT(ncols % GGML_SYCL_DMMV_X == 0);
-    const int block_num_y = (nrows + GGML_SYCL_MMV_Y - 1) / GGML_SYCL_MMV_Y;
-    const sycl::range<3> block_nums(1, 1, block_num_y);
-    const sycl::range<3> block_dims(1, GGML_SYCL_MMV_Y, WARP_SIZE);
-    {
-        dpct::has_capability_or_fail(stream->get_device(),
-                                     {sycl::aspect::fp16});
-
-        stream->parallel_for(
-            sycl::nd_range<3>(block_nums * block_dims, block_dims),
-            [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                dequantize_mul_mat_vec<QK5_0, QR5_0, dequantize_q5_0>(
-                    vx, y, dst, ncols, nrows, item_ct1);
-            });
-    }
-}
-
-static void dequantize_mul_mat_vec_q5_1_sycl(const void *vx, const dfloat *y,
-                                             float *dst, const int ncols,
-                                             const int nrows,
-                                             dpct::queue_ptr stream) {
-    GGML_ASSERT(ncols % GGML_SYCL_DMMV_X == 0);
-    const int block_num_y = (nrows + GGML_SYCL_MMV_Y - 1) / GGML_SYCL_MMV_Y;
-    const sycl::range<3> block_nums(1, 1, block_num_y);
-    const sycl::range<3> block_dims(1, GGML_SYCL_MMV_Y, WARP_SIZE);
-    {
-        dpct::has_capability_or_fail(stream->get_device(),
-                                     {sycl::aspect::fp16});
-
-        stream->parallel_for(
-            sycl::nd_range<3>(block_nums * block_dims, block_dims),
-            [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                dequantize_mul_mat_vec<QK5_1, QR5_1, dequantize_q5_1>(
-                    vx, y, dst, ncols, nrows, item_ct1);
-            });
-    }
-}
-
-static void dequantize_mul_mat_vec_q8_0_sycl_reorder(const void *vx, const dfloat *y,
-                                             float *dst, const int ncols,
-                                             const int nrows,
-                                             dpct::queue_ptr stream) {
-    GGML_ASSERT(ncols % GGML_SYCL_DMMV_X == 0);
-    const int block_num_y = (nrows + GGML_SYCL_MMV_Y - 1) / GGML_SYCL_MMV_Y;
-    const sycl::range<3> block_nums(1, 1, block_num_y);
-    const sycl::range<3> block_dims(1, GGML_SYCL_MMV_Y, WARP_SIZE);
-    {
-        dpct::has_capability_or_fail(stream->get_device(),
-                                     {sycl::aspect::fp16});
-
-        stream->parallel_for(
-            sycl::nd_range<3>(block_nums * block_dims, block_dims),
-            [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                // Q8_0 reorder layout: [all qs (ncols*nrows bytes)][all d values]
-                // Cannot reuse dequantize_mul_mat_vec_reorder template because it has
-                // Q4_0-specific constants hardcoded (d_ptr offset and qs stride).
-                const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
-                    item_ct1.get_local_id(1);
-                if (row >= nrows) return;
-
-                const int tid = item_ct1.get_local_id(2);
-                const int iter_stride = 8*2*GGML_SYCL_DMMV_X;
-                const int vals_per_iter = iter_stride / WARP_SIZE;
-                const int ncols_left = ncols % (QK8_0*WARP_SIZE);
-                const int ncols_align = ncols - ncols_left;
-
-#ifdef GGML_SYCL_F16
-                sycl::half2 tmp = {0.0f, 0.0f};
-#else
-                float tmp = 0.0f;
-#endif
-                const char *d_ptr = (const char*)vx + ncols*nrows;  // d after all qs
-
-                int i = 0;
-                for (i = 0; i < ncols_align; i += iter_stride) {
-                    const int col = i + vals_per_iter*tid;
-                    const int ib = (row*ncols + col)/QK8_0;
-                    const int iqs = col % QK8_0;
-
-#pragma unroll
-                    for (int j = 0; j < vals_per_iter; j += 2) {
-                        dfloat2 v;
-                        dequantize_q8_0_reorder((const void *)d_ptr, ib, (const void *)vx,
-                                                ib * QK8_0 + iqs + j, v);
-
-#ifdef GGML_SYCL_F16
-                        dfloat2 t1{y[col + j + 0], y[col + j + 1]};
-                        tmp += v * t1;
-#else
-                        tmp += v.x() * y[col + j + 0];
-                        tmp += v.y() * y[col + j + 1];
-#endif
-                    }
-                }
-
-                // handle remaining columns
-                for (; i < ncols; i += iter_stride) {
-                    if (tid >= ncols_left/QK8_0) continue;
-                    const int col = i + vals_per_iter*tid;
-                    const int ib = (row*ncols + col)/QK8_0;
-                    const int iqs = col % QK8_0;
-
-#pragma unroll
-                    for (int j = 0; j < vals_per_iter; j += 2) {
-                        dfloat2 v;
-                        dequantize_q8_0_reorder((const void *)d_ptr, ib, (const void *)vx,
-                                                ib * QK8_0 + iqs + j, v);
-
-#ifdef GGML_SYCL_F16
-                        dfloat2 t1{y[col + j + 0], y[col + j + 1]};
-                        tmp += v * t1;
-#else
-                        tmp += v.x() * y[col + j + 0];
-                        tmp += v.y() * y[col + j + 1];
-#endif
-                    }
-                }
-
-                // reduce
-                const int mask_start = ncols > GGML_SYCL_DMMV_X ? WARP_SIZE >> 1 : WARP_SIZE >> 2;
-                for (int mask = mask_start; mask > 0; mask >>= 1) {
-                    tmp += dpct::permute_sub_group_by_xor(item_ct1.get_sub_group(), tmp, mask);
-                }
-
-                if (tid == 0) {
-#ifdef GGML_SYCL_F16
-                    dst[row] = tmp.x() + tmp.y();
-#else
-                    dst[row] = tmp;
-#endif
-                }
-            });
-    }
-}
-
-static void dequantize_mul_mat_vec_q8_0_sycl(const void *vx, const dfloat *y,
-                                             float *dst, const int ncols,
-                                             const int nrows,
-                                             dpct::queue_ptr stream) {
-    GGML_ASSERT(ncols % GGML_SYCL_DMMV_X == 0);
-    const int block_num_y = (nrows + GGML_SYCL_MMV_Y - 1) / GGML_SYCL_MMV_Y;
-    const sycl::range<3> block_nums(1, 1, block_num_y);
-    const sycl::range<3> block_dims(1, GGML_SYCL_MMV_Y, WARP_SIZE);
-    {
-        dpct::has_capability_or_fail(stream->get_device(),
-                                     {sycl::aspect::fp16});
-
-        stream->parallel_for(
-            sycl::nd_range<3>(block_nums * block_dims, block_dims),
-            [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                dequantize_mul_mat_vec<QK8_0, QR8_0, dequantize_q8_0>(
-                    vx, y, dst, ncols, nrows, item_ct1);
-            });
-    }
-}
-
-static void dequantize_mul_mat_vec_q2_K_sycl(const void *vx, const float *y,
-                                             float *dst, const int ncols,
-                                             const int nrows,
-                                             dpct::queue_ptr stream) {
+template <int ny, auto kernel>
+static void launch_dequantize_mul_mat_vec_k(const void *vx, const float *y, float *dst,
+                                            const int ncols, const int nrows,
+                                            dpct::queue_ptr stream) {
     GGML_ASSERT(ncols % QK_K == 0);
-    const int ny = 2; // very slightly faster than 1 even when K_QUANTS_PER_ITERATION = 2
     const int block_num_y = (nrows + ny - 1) / ny;
     const sycl::range<3> block_nums(1, 1, block_num_y);
     const sycl::range<3> block_dims(1, ny, WARP_SIZE);
     stream->parallel_for(
         sycl::nd_range<3>(block_nums * block_dims, block_dims),
         [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-            dequantize_mul_mat_vec_q2_k(vx, y, dst, ncols, nrows, item_ct1);
-        });
-}
-
-static void dequantize_mul_mat_vec_q2_K_sycl_reorder(const void *vx, const float *y,
-                                                     float *dst, const int ncols,
-                                                     const int nrows,
-                                                     dpct::queue_ptr stream) {
-    GGML_ASSERT(ncols % QK_K == 0);
-    const int ny = 2 / K_QUANTS_PER_ITERATION;
-    const int block_num_y = (nrows + ny - 1) / ny;
-    const sycl::range<3> block_nums(1, 1, block_num_y);
-    const sycl::range<3> block_dims(1, ny, WARP_SIZE);
-    stream->parallel_for(
-        sycl::nd_range<3>(block_nums * block_dims, block_dims),
-        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-            dequantize_mul_mat_vec_q2_k_reorder(vx, y, dst, ncols, nrows, item_ct1);
-        });
-}
-
-static void dequantize_mul_mat_vec_q3_K_sycl(const void *vx, const float *y,
-                                             float *dst, const int ncols,
-                                             const int nrows,
-                                             dpct::queue_ptr stream) {
-    GGML_ASSERT(ncols % QK_K == 0);
-    const int ny = 2 / K_QUANTS_PER_ITERATION;
-    const int block_num_y = (nrows + ny - 1) / ny;
-    const sycl::range<3> block_nums(1, 1, block_num_y);
-    const sycl::range<3> block_dims(1, ny, WARP_SIZE);
-    stream->parallel_for(
-        sycl::nd_range<3>(block_nums * block_dims, block_dims),
-        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-            dequantize_mul_mat_vec_q3_k(vx, y, dst, ncols, nrows, item_ct1);
-        });
-}
-
-static void dequantize_mul_mat_vec_q3_K_sycl_reorder(const void *vx, const float *y,
-                                                     float *dst, const int ncols,
-                                                     const int nrows,
-                                                     dpct::queue_ptr stream) {
-    GGML_ASSERT(ncols % QK_K == 0);
-    const int ny = 2 / K_QUANTS_PER_ITERATION;
-    const int block_num_y = (nrows + ny - 1) / ny;
-    const sycl::range<3> block_nums(1, 1, block_num_y);
-    const sycl::range<3> block_dims(1, ny, WARP_SIZE);
-    stream->parallel_for(
-        sycl::nd_range<3>(block_nums * block_dims, block_dims),
-        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-            dequantize_mul_mat_vec_q3_k_reorder(vx, y, dst, ncols, nrows, item_ct1);
-        });
-}
-
-static void dequantize_mul_mat_vec_q4_K_sycl(const void *vx, const float *y,
-                                             float *dst, const int ncols,
-                                             const int nrows,
-                                             dpct::queue_ptr stream) {
-    GGML_ASSERT(ncols % QK_K == 0);
-    const int ny = 2 / K_QUANTS_PER_ITERATION;
-    const int block_num_y = (nrows + ny - 1) / ny;
-    const sycl::range<3> block_nums(1, 1, block_num_y);
-    const sycl::range<3> block_dims(1, ny, WARP_SIZE);
-    stream->parallel_for(
-        sycl::nd_range<3>(block_nums * block_dims, block_dims),
-        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-            dequantize_mul_mat_vec_q4_k(vx, y, dst, ncols, nrows, item_ct1);
-        });
-}
-
-static void dequantize_mul_mat_vec_q5_K_sycl(const void *vx, const float *y,
-                                             float *dst, const int ncols,
-                                             const int nrows,
-                                             dpct::queue_ptr stream) {
-    GGML_ASSERT(ncols % QK_K == 0);
-    const int ny = 2 / K_QUANTS_PER_ITERATION;
-    const int block_num_y = (nrows + ny - 1) / ny;
-    const sycl::range<3> block_nums(1, 1, block_num_y);
-    const sycl::range<3> block_dims(1, ny, WARP_SIZE);
-    stream->parallel_for(
-        sycl::nd_range<3>(block_nums * block_dims, block_dims),
-        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-            dequantize_mul_mat_vec_q5_k(vx, y, dst, ncols, nrows, item_ct1);
-        });
-}
-
-static void dequantize_mul_mat_vec_q6_K_sycl(const void *vx, const float *y,
-                                             float *dst, const int ncols,
-                                             const int nrows,
-                                             dpct::queue_ptr stream) {
-    GGML_ASSERT(ncols % QK_K == 0);
-    const int ny = 2 / K_QUANTS_PER_ITERATION;
-    const int block_num_y = (nrows + ny - 1) / ny;
-    const sycl::range<3> block_nums(1, 1, block_num_y);
-    const sycl::range<3> block_dims(1, ny, WARP_SIZE);
-    stream->parallel_for(
-        sycl::nd_range<3>(block_nums * block_dims, block_dims),
-        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-            dequantize_mul_mat_vec_q6_k(vx, y, dst, ncols, nrows, item_ct1);
+            kernel(vx, y, dst, ncols, nrows, item_ct1);
         });
 }
 
@@ -2248,54 +1926,6 @@ bool ggml_sycl_dequantize_mul_mat_vec_glu_reorder_esimd(enum ggml_type src0_type
 
 #endif // GGML_SYCL_DMMV_HAS_ESIMD
 
-static void dequantize_mul_mat_vec_q4_K_sycl_reorder(const void *vx, const float *y,
-                                                     float *dst, const int ncols,
-                                                     const int nrows,
-                                                     dpct::queue_ptr stream) {
-    GGML_ASSERT(ncols % QK_K == 0);
-    const int ny = 2 / K_QUANTS_PER_ITERATION;
-    const int block_num_y = (nrows + ny - 1) / ny;
-    const sycl::range<3> block_nums(1, 1, block_num_y);
-    const sycl::range<3> block_dims(1, ny, WARP_SIZE);
-    stream->parallel_for(
-        sycl::nd_range<3>(block_nums * block_dims, block_dims),
-        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-            dequantize_mul_mat_vec_q4_k_reorder(vx, y, dst, ncols, nrows, item_ct1);
-        });
-}
-
-static void dequantize_mul_mat_vec_q5_K_sycl_reorder(const void *vx, const float *y,
-                                                     float *dst, const int ncols,
-                                                     const int nrows,
-                                                     dpct::queue_ptr stream) {
-    GGML_ASSERT(ncols % QK_K == 0);
-    const int ny = 2 / K_QUANTS_PER_ITERATION;
-    const int block_num_y = (nrows + ny - 1) / ny;
-    const sycl::range<3> block_nums(1, 1, block_num_y);
-    const sycl::range<3> block_dims(1, ny, WARP_SIZE);
-    stream->parallel_for(
-        sycl::nd_range<3>(block_nums * block_dims, block_dims),
-        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-            dequantize_mul_mat_vec_q5_k_reorder(vx, y, dst, ncols, nrows, item_ct1);
-        });
-}
-
-static void dequantize_mul_mat_vec_q6_K_sycl_reorder(const void *vx, const float *y,
-                                                     float *dst, const int ncols,
-                                                     const int nrows,
-                                                     dpct::queue_ptr stream) {
-    GGML_ASSERT(ncols % QK_K == 0);
-    const int ny = 2 / K_QUANTS_PER_ITERATION;
-    const int block_num_y = (nrows + ny - 1) / ny;
-    const sycl::range<3> block_nums(1, 1, block_num_y);
-    const sycl::range<3> block_dims(1, ny, WARP_SIZE);
-    stream->parallel_for(
-        sycl::nd_range<3>(block_nums * block_dims, block_dims),
-        [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-            dequantize_mul_mat_vec_q6_k_reorder(vx, y, dst, ncols, nrows, item_ct1);
-        });
-}
-
 void ggml_sycl_op_dequantize_mul_mat_vec(
     ggml_backend_sycl_context & ctx,
     const ggml_tensor *src0, const ggml_tensor *src1, ggml_tensor *dst,
@@ -2344,27 +1974,27 @@ void ggml_sycl_op_dequantize_mul_mat_vec(
         case GGML_TYPE_Q1_0:
             if ((ggml_tensor_extra_gpu*)dst->src[0]->extra &&
                 ((ggml_tensor_extra_gpu*)dst->src[0]->extra)->optimized_feature.reorder) {
-                dequantize_mul_mat_vec_q1_0_sycl_reorder(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
+                launch_dequantize_mul_mat_vec_reorder<QK1_0, QR1_0, QK1_0/8, dequantize_q1_0_reorder>(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
             } else {
-                dequantize_mul_mat_vec_q1_0_sycl(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
+                launch_dequantize_mul_mat_vec<QK1_0, QR1_0, dequantize_q1_0>(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
             }
             break;
         case GGML_TYPE_Q4_0:
             if ((ggml_tensor_extra_gpu*)dst->src[0]->extra &&
                 ((ggml_tensor_extra_gpu*)dst->src[0]->extra)->optimized_feature.reorder) {
-                dequantize_mul_mat_vec_q4_0_sycl_reorder(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
+                launch_dequantize_mul_mat_vec_reorder<QK4_0, QR4_0, QK4_0/2, dequantize_q4_0_reorder>(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
             } else {
-                dequantize_mul_mat_vec_q4_0_sycl(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
+                launch_dequantize_mul_mat_vec<QK4_0, QR4_0, dequantize_q4_0>(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
             }
             break;
         case GGML_TYPE_Q4_1:
-            dequantize_mul_mat_vec_q4_1_sycl(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
+            launch_dequantize_mul_mat_vec<QK4_1, QR4_1, dequantize_q4_1>(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
             break;
         case GGML_TYPE_Q5_0:
-            dequantize_mul_mat_vec_q5_0_sycl(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
+            launch_dequantize_mul_mat_vec<QK5_0, QR5_0, dequantize_q5_0>(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
             break;
         case GGML_TYPE_Q5_1:
-            dequantize_mul_mat_vec_q5_1_sycl(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
+            launch_dequantize_mul_mat_vec<QK5_1, QR5_1, dequantize_q5_1>(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
             break;
         case GGML_TYPE_Q8_0:
             if ((ggml_tensor_extra_gpu *) dst->src[0]->extra &&
@@ -2375,10 +2005,10 @@ void ggml_sycl_op_dequantize_mul_mat_vec(
                 } else
 #endif
                 {
-                    dequantize_mul_mat_vec_q8_0_sycl_reorder(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
+                    launch_dequantize_mul_mat_vec_reorder<QK8_0, QR8_0, QK8_0, dequantize_q8_0_reorder>(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
                 }
             } else {
-                dequantize_mul_mat_vec_q8_0_sycl(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
+                launch_dequantize_mul_mat_vec<QK8_0, QR8_0, dequantize_q8_0>(src0_dd_i, src1_dfloat, dst_dd_i, ne00, row_diff, stream);
             }
             break;
         case GGML_TYPE_Q2_K:
@@ -2391,10 +2021,10 @@ void ggml_sycl_op_dequantize_mul_mat_vec(
                 else
 #endif
                 {
-                    dequantize_mul_mat_vec_q2_K_sycl_reorder(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
+                    launch_dequantize_mul_mat_vec_k<2/K_QUANTS_PER_ITERATION, dequantize_mul_mat_vec_q2_k_reorder>(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
                 }
             } else {
-                dequantize_mul_mat_vec_q2_K_sycl(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
+                launch_dequantize_mul_mat_vec_k<2, dequantize_mul_mat_vec_q2_k>(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
             }
             break;
         case GGML_TYPE_Q3_K:
@@ -2407,10 +2037,10 @@ void ggml_sycl_op_dequantize_mul_mat_vec(
                 else
 #endif
                 {
-                    dequantize_mul_mat_vec_q3_K_sycl_reorder(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
+                    launch_dequantize_mul_mat_vec_k<2/K_QUANTS_PER_ITERATION, dequantize_mul_mat_vec_q3_k_reorder>(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
                 }
             } else {
-                dequantize_mul_mat_vec_q3_K_sycl(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
+                launch_dequantize_mul_mat_vec_k<2/K_QUANTS_PER_ITERATION, dequantize_mul_mat_vec_q3_k>(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
             }
             break;
         case GGML_TYPE_Q4_K:
@@ -2423,10 +2053,10 @@ void ggml_sycl_op_dequantize_mul_mat_vec(
                 else
 #endif
                 {
-                    dequantize_mul_mat_vec_q4_K_sycl_reorder(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
+                    launch_dequantize_mul_mat_vec_k<2/K_QUANTS_PER_ITERATION, dequantize_mul_mat_vec_q4_k_reorder>(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
                 }
             } else {
-                dequantize_mul_mat_vec_q4_K_sycl(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
+                launch_dequantize_mul_mat_vec_k<2/K_QUANTS_PER_ITERATION, dequantize_mul_mat_vec_q4_k>(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
             }
             break;
         case GGML_TYPE_Q5_K:
@@ -2439,10 +2069,10 @@ void ggml_sycl_op_dequantize_mul_mat_vec(
                 else
 #endif
                 {
-                    dequantize_mul_mat_vec_q5_K_sycl_reorder(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
+                    launch_dequantize_mul_mat_vec_k<2/K_QUANTS_PER_ITERATION, dequantize_mul_mat_vec_q5_k_reorder>(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
                 }
             } else {
-                dequantize_mul_mat_vec_q5_K_sycl(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
+                launch_dequantize_mul_mat_vec_k<2/K_QUANTS_PER_ITERATION, dequantize_mul_mat_vec_q5_k>(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
             }
             break;
         case GGML_TYPE_Q6_K:
@@ -2455,10 +2085,10 @@ void ggml_sycl_op_dequantize_mul_mat_vec(
                 else
 #endif
                 {
-                    dequantize_mul_mat_vec_q6_K_sycl_reorder(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
+                    launch_dequantize_mul_mat_vec_k<2/K_QUANTS_PER_ITERATION, dequantize_mul_mat_vec_q6_k_reorder>(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
                 }
             } else {
-                dequantize_mul_mat_vec_q6_K_sycl(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
+                launch_dequantize_mul_mat_vec_k<2/K_QUANTS_PER_ITERATION, dequantize_mul_mat_vec_q6_k>(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream);
             }
             break;
         case GGML_TYPE_F16:

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

@@ -7653,7 +7653,8 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
         case GGML_OP_SOLVE_TRI:
             return op->src[0]->ne[0] <= SYCL_SOLVE_TRI_MAX_N && op->src[1]->ne[0] <= SYCL_SOLVE_TRI_MAX_K;
         case GGML_OP_FLASH_ATTN_EXT:
-            return ggml_sycl_flash_attn_ext_supported(device, op);
+            // K/V row lists (ggml_flash_attn_ext_rows) are not supported
+            return op->src[5] == nullptr && ggml_sycl_flash_attn_ext_supported(device, op);
         default:
             return false;
     }
@@ -7829,9 +7830,15 @@ struct ggml_backend_sycl_comm_context {
 
     // pinned host staging per device context, reused across calls (queues are in-order)
     uint8_t * host_out[2] = { nullptr, nullptr };
-    uint8_t * host_in[2]  = { nullptr, nullptr };
+    uint8_t * host_in[2]  = { nullptr, nullptr };  // only used by the 4-buffer path (separate SYCL contexts)
     size_t    host_bytes  = 0;
 
+    // SYCL_ENABLE_DEFAULT_CONTEXTS=1 sets comm_exchange_fast on Linux - PR #30956
+    using exchange_fn_t = void (*)(ggml_backend_sycl_comm_context * comm, queue_ptr q0, queue_ptr q1,
+                                   void * dst0, const void * src1, void * dst1, const void * src0,
+                                   size_t nbytes);
+    exchange_fn_t comm_exchange = nullptr; // Moves same context check out of hot path; tests only once
+
     void free_host() {
         for (int i = 0; i < 2; ++i) {
             const sycl::context sctx = ((ggml_backend_sycl_context *) backends[i]->context)->stream()->get_context();
@@ -7842,6 +7849,41 @@ struct ggml_backend_sycl_comm_context {
     }
 };
 
+// dst0 <- src1 and dst1 <- src0 through pinned host memory.
+static void ggml_sycl_comm_exchange_fast(ggml_backend_sycl_comm_context * comm, queue_ptr q0, queue_ptr q1,
+                                         void * dst0, const void * src1, void * dst1, const void * src0,
+                                         size_t nbytes) {
+
+    // Increases Linux tensor split prompt eval speeds on par with layer split speeds
+    // Reallocate host buffers if they are too small
+    if (comm->host_bytes < nbytes) {
+        q0->wait();
+        q1->wait();
+        comm->free_host();
+        comm->host_out[0] = sycl::malloc_host<uint8_t>(nbytes, q0->get_context());
+        comm->host_out[1] = sycl::malloc_host<uint8_t>(nbytes, q1->get_context());
+
+        GGML_ASSERT(comm->host_out[0] && comm->host_out[1]);
+        comm->host_bytes = nbytes;
+    }
+
+    // Step 1: Concurrently copy from GPUs into host buffers (Asynchronous)
+    sycl::event e_d2h_0 = q0->memcpy(comm->host_out[0], src0, nbytes);
+    sycl::event e_d2h_1 = q1->memcpy(comm->host_out[1], src1, nbytes);
+
+    // Step 2: Asynchronous Barrier Interlocking (NO HOST BLOCKING), wait() causes MTP conflicts
+    q0->ext_oneapi_submit_barrier({e_d2h_1}); // q0 waits for q1's device-to-host copy
+    q1->ext_oneapi_submit_barrier({e_d2h_0}); // q1 waits for q0's device-to-host copy
+
+    // Step 3: Cross-upload directly from the opposing host buffer.
+    sycl::event e_h2d_0 = q0->memcpy(dst0, comm->host_out[1], nbytes); // q0 pulls from host_out[1]
+    sycl::event e_h2d_1 = q1->memcpy(dst1, comm->host_out[0], nbytes); // q1 pulls from host_out[0]
+
+    // Step 4: Finalize the exchange before returning to host code
+    e_h2d_0.wait();
+    e_h2d_1.wait();
+}
+
 // dst0 <- src1 and dst1 <- src0 through pinned host memory, with a single host wait
 static void ggml_sycl_comm_exchange(ggml_backend_sycl_comm_context * comm, queue_ptr q0, queue_ptr q1,
                                     void * dst0, const void * src1, void * dst1, const void * src0,
@@ -7887,6 +7929,15 @@ void * ggml_backend_sycl_comm_init(ggml_backend_t * backends, size_t n_backends)
     auto * sctx1 = (ggml_backend_sycl_context *) backends[1]->context;
     ctx->buf0 = std::make_unique<ggml_sycl_pool_alloc<uint8_t>>(sctx0->pool());
     ctx->buf1 = std::make_unique<ggml_sycl_pool_alloc<uint8_t>>(sctx1->pool());
+
+#if defined(_WIN32) || defined(_WIN64) // Not supported on Windows
+    ctx->comm_exchange = ggml_sycl_comm_exchange;
+#else
+    // SYCL_ENABLE_DEFAULT_CONTEXTS=1 sets comm_exchange_fast on Linux - PR #30956
+    ctx->comm_exchange = (sctx0->stream()->get_context() == sctx1->stream()->get_context())
+                       ? ggml_sycl_comm_exchange_fast
+                       : ggml_sycl_comm_exchange;
+#endif
     return ctx;
 }
 catch (const sycl::exception &) { return nullptr; }
@@ -7970,7 +8021,7 @@ bool ggml_backend_sycl_comm_allreduce_tensor(void * comm_ctx_v, struct ggml_tens
         sycl::half * f16_tmp0 = (sycl::half *) buf0;
         sycl::half * f16_tmp1 = (sycl::half *) buf1;
 
-        ggml_sycl_comm_exchange(comm_ctx, q0, q1, f16_tmp0, tensors[1]->data, f16_tmp1, tensors[0]->data, nbytes);
+        comm_ctx->comm_exchange(comm_ctx, q0, q1, f16_tmp0, tensors[1]->data, f16_tmp1, tensors[0]->data, nbytes);
 
         q0->submit([&](sycl::handler & h) {
             h.parallel_for(sycl::range<1>(nelem), [=](sycl::id<1> i) {
@@ -8000,7 +8051,7 @@ bool ggml_backend_sycl_comm_allreduce_tensor(void * comm_ctx_v, struct ggml_tens
         float * tmp1 = (float *) buf1;
 
         // separate SYCL contexts: a raw peer-USM q->memcpy is a silent no-op, so stage through host memory
-        ggml_sycl_comm_exchange(comm_ctx, q0, q1, tmp0, tensors[1]->data, tmp1, tensors[0]->data, nbytes);
+        comm_ctx->comm_exchange(comm_ctx, q0, q1, tmp0, tensors[1]->data, tmp1, tensors[0]->data, nbytes);
 
         q0->submit([&](sycl::handler & h) {
             h.parallel_for(sycl::range<1>(nelem), [=](sycl::id<1> i) {
@@ -8034,7 +8085,7 @@ bool ggml_backend_sycl_comm_allreduce_tensor(void * comm_ctx_v, struct ggml_tens
 
     // Phase B: exchange the compressed bytes (in-order queues keep this after the compress kernels).
     const size_t bf16_bytes = nelem * sizeof(uint16_t);
-    ggml_sycl_comm_exchange(comm_ctx, q0, q1, inbox0, outbox1, inbox1, outbox0, bf16_bytes);
+    comm_ctx->comm_exchange(comm_ctx, q0, q1, inbox0, outbox1, inbox1, outbox0, bf16_bytes);
 
     // Phase C: decompress + add into local FP32 partial.
     q0->submit([&](sycl::handler & h) {

ggml/src/ggml-vulkan/CMakeLists.txt

@@ -66,6 +66,7 @@ if (Vulkan_FOUND)
                              ggml-vulkan-push-constants.h
                              ggml-vulkan-common.h
                              ggml-vulkan-buffers.cpp
+                             ggml-vulkan-comm.cpp
                              ggml-vulkan-debug.cpp
                             )
 

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

@@ -0,0 +1,591 @@
+#include "ggml-vulkan-common.h"
+
+#if !defined(_WIN32)
+#include <fcntl.h>
+#include <unistd.h>
+#if defined(__linux__)
+#include <pthread.h>
+#endif
+
+static constexpr size_t ggml_vk_comm_f16_min_size = 1024 * 1024;
+
+// follows ops ggml-backend-meta may delay a reduction across
+// TODO: the required precision should come from the meta backend instead.
+static bool ggml_vk_comm_needs_f32(const ggml_tensor * t, int depth = 0) {
+    if (t == nullptr || depth > 64) {
+        return false;
+    }
+    switch (t->op) {
+        case GGML_OP_MUL_MAT:
+        case GGML_OP_MUL_MAT_ID:
+            return ggml_get_op_params_i32(t, 0) == GGML_PREC_F32;
+        case GGML_OP_ADD:
+            return ggml_vk_comm_needs_f32(t->src[0], depth + 1) || ggml_vk_comm_needs_f32(t->src[1], depth + 1);
+        case GGML_OP_MUL:
+        case GGML_OP_ADD_ID:
+            return ggml_vk_comm_needs_f32(t->src[0], depth + 1);
+        case GGML_OP_VIEW:
+            return ggml_vk_comm_needs_f32(t->src[0], depth + 1);
+        default:
+            return false;
+    }
+}
+
+// host_bridge: device i signals done[i], a host thread signals go[i] once all peers of i are done
+// sync_fd: a host thread exchanges SYNC_FD semaphores between peers, the wait resolves in the kernel
+enum class vk_comm_sync {
+    host_bridge,
+    sync_fd,
+};
+
+struct vk_comm_device {
+    ggml_backend_t                     backend  = nullptr;
+    ggml_backend_vk_context *          ctx      = nullptr;
+    void *                             host_ptr = nullptr;
+    std::vector<ggml_backend_buffer_t> host_buf; // host_ptr imported into each peer
+
+    struct {
+        PFN_vkGetSemaphoreFdKHR    get_fd    = nullptr;
+        PFN_vkImportSemaphoreFdKHR import_fd = nullptr;
+        vk::Semaphore              sem;
+        std::vector<vk::Semaphore> imp;      // [peer], pooled, reused only after ggml_vk_graph_cleanup
+    } sfd;
+
+    struct {
+        vk::Semaphore done;
+        vk::Semaphore go;
+    } bridge;
+};
+
+struct ggml_backend_vk_comm_context {
+    std::vector<vk_comm_device> devs;
+    vk_comm_sync                mode     = vk_comm_sync::host_bridge;
+    size_t                      align    = 4096;
+    size_t                      cap      = 0;
+    uint64_t                    round    = 0;
+    ggml_context *              tctx     = nullptr;
+    ggml_tensor *               slot     = nullptr;
+    bool                        disabled = false;
+
+    std::thread                 thread;
+    std::chrono::nanoseconds    spin{ 200'000 };
+    std::atomic<uint64_t>       submitted{ 0 };
+    std::atomic<uint64_t>       imported{ 0 };
+    std::atomic<bool>           sleeping{ false };
+    std::atomic<bool>           stop{ false };
+    std::mutex                  mutex;
+    std::condition_variable     cv;
+};
+
+static void ggml_vk_comm_free_buffers(ggml_backend_vk_comm_context * comm) {
+    for (vk_comm_device & d : comm->devs) {
+        for (ggml_backend_buffer_t & buf : d.host_buf) {
+            ggml_backend_buffer_free(buf);
+            buf = nullptr;
+        }
+        std::free(d.host_ptr);
+        d.host_ptr = nullptr;
+    }
+    comm->cap = 0;
+}
+
+static bool ggml_backend_vk_comm_ensure(ggml_backend_vk_comm_context * comm, size_t nbytes) {
+    if (comm->disabled) {
+        return false;
+    }
+    if (nbytes <= comm->cap) {
+        return true;
+    }
+    auto fail = [comm]() {
+        GGML_LOG_WARN("ggml_vulkan: shared host buffer allocation failed, using generic all-reduce\n");
+        comm->disabled = true;
+        return false;
+    };
+    for (vk_comm_device & d : comm->devs) {
+        ggml_backend_synchronize(d.backend);
+    }
+    ggml_vk_comm_free_buffers(comm);
+
+    const size_t n   = comm->devs.size();
+    const size_t cap = GGML_PAD(nbytes, comm->align);
+    for (size_t i = 0; i < n; i++) {
+        vk_comm_device & d = comm->devs[i];
+        if (posix_memalign(&d.host_ptr, comm->align, 2 * cap) != 0) {
+            d.host_ptr = nullptr;
+            return fail();
+        }
+        for (size_t j = 0; j < n; j++) {
+            d.host_buf[j] = ggml_backend_vk_device_buffer_from_host_ptr(
+                ggml_backend_get_device(comm->devs[j].backend), d.host_ptr, 2 * cap, 2 * cap);
+            if (!d.host_buf[j]) {
+                return fail();
+            }
+        }
+    }
+    comm->cap = cap;
+    return true;
+}
+
+static void ggml_vk_comm_set_slot(ggml_tensor * slot, const ggml_tensor * src, ggml_type type, ggml_backend_buffer_t buf, size_t offs) {
+    slot->type  = type;
+    slot->nb[0] = ggml_type_size(type);
+    for (int d = 0; d < GGML_MAX_DIMS; d++) {
+        slot->ne[d] = src->ne[d];
+        if (d > 0) {
+            slot->nb[d] = slot->nb[d - 1] * slot->ne[d - 1];
+        }
+    }
+    slot->buffer = buf;
+    slot->data   = (char *) ggml_backend_buffer_get_base(buf) + offs;
+}
+
+static void ggml_vk_comm_submit(ggml_backend_vk_context * ctx, vk_context & c, vk::Semaphore signal, uint64_t value) {
+    ggml_vk_ctx_end(c);
+    for (auto & cpy : c->in_memcpys) {
+        memcpy(cpy.dst, cpy.src, cpy.n);
+    }
+    for (auto & mset : c->memsets) {
+        memset(mset.dst, mset.val, mset.n);
+    }
+    c->in_memcpys.clear();
+    c->memsets.clear();
+    if (signal) {
+        c->seqs.back().back().signal_semaphores.push_back({ signal, value });
+    }
+    ggml_vk_submit(c, {});
+    ctx->submit_pending = true;
+    ctx->compute_ctx.reset();
+}
+
+static void ggml_vk_comm_notify_submitted(ggml_backend_vk_comm_context * comm, uint64_t round) {
+    comm->submitted.store(round);
+    if (comm->sleeping.load()) {
+        std::lock_guard<std::mutex> lock(comm->mutex);
+        comm->cv.notify_one();
+    }
+}
+
+// false on stop
+static bool ggml_vk_comm_wait_submitted(ggml_backend_vk_comm_context * comm, uint64_t r) {
+    using clock = std::chrono::steady_clock;
+    for (const auto t0 = clock::now(); comm->submitted.load() < r;) {
+        if (comm->stop.load()) {
+            return false;
+        }
+        if (clock::now() - t0 < comm->spin) {
+            YIELD();
+            continue;
+        }
+        comm->sleeping.store(true);
+        if (comm->submitted.load() < r) {
+            std::unique_lock<std::mutex> lock(comm->mutex);
+            comm->cv.wait(lock, [&] { return comm->submitted.load() >= r || comm->stop.load(); });
+        }
+        comm->sleeping.store(false);
+    }
+    return true;
+}
+
+// signals go[i] once all peers of i are done; false on stop
+static bool ggml_vk_comm_bridge_release(ggml_backend_vk_comm_context * comm, uint64_t r, std::vector<uint64_t> & seen, std::vector<bool> & sent) {
+    using clock = std::chrono::steady_clock;
+    const size_t n = comm->devs.size();
+    std::fill(sent.begin(), sent.end(), false);
+    size_t left = n;
+    for (const auto t0 = clock::now();;) {
+        size_t pending = n;
+        for (size_t i = 0; i < n; i++) {
+            if (sent[i]) {
+                continue;
+            }
+            bool ready = true;
+            for (size_t j = 0; j < n && ready; j++) {
+                if (j != i && seen[j] < r) {
+                    const vk_comm_device & p = comm->devs[j];
+                    seen[j] = p.ctx->device->device.getSemaphoreCounterValue(p.bridge.done);
+                    ready   = seen[j] >= r;
+                    pending = ready ? pending : j;
+                }
+            }
+            if (ready) {
+                const vk_comm_device & d = comm->devs[i];
+                d.ctx->device->device.signalSemaphore({ d.bridge.go, r });
+                sent[i] = true;
+                left--;
+            }
+        }
+        if (left == 0) {
+            return true;
+        }
+        if (clock::now() - t0 < comm->spin) {
+            YIELD();
+            continue;
+        }
+        const vk_comm_device &      p = comm->devs[pending];
+        const vk::SemaphoreWaitInfo wi{ {}, 1, &p.bridge.done, &r };
+        while (p.ctx->device->device.waitSemaphores(wi, 100'000'000) != vk::Result::eSuccess) {
+            if (comm->stop.load()) {
+                return false;
+            }
+        }
+    }
+}
+
+static void ggml_vk_comm_bridge_thread(ggml_backend_vk_comm_context * comm) {
+#if defined(__linux__)
+    pthread_setname_np(pthread_self(), "ggml-vk-comm");
+#endif
+    std::vector<uint64_t> seen(comm->devs.size(), 0);
+    std::vector<bool>     sent(comm->devs.size());
+    try {
+        for (uint64_t r = 1;; r++) {
+            if (!ggml_vk_comm_wait_submitted(comm, r) || !ggml_vk_comm_bridge_release(comm, r, seen, sent)) {
+                return;
+            }
+        }
+    } catch (const vk::SystemError & e) {
+        GGML_ABORT("ggml_vulkan: all-reduce bridge failed: %s", e.what());
+    }
+}
+
+static void ggml_vk_comm_bridge_init(ggml_backend_vk_comm_context * comm) {
+    vk::SemaphoreTypeCreateInfo stci{ vk::SemaphoreType::eTimeline, 0 };
+    vk::SemaphoreCreateInfo     sci{};
+    sci.pNext = &stci;
+    for (vk_comm_device & d : comm->devs) {
+        d.bridge.done = d.ctx->device->device.createSemaphore(sci);
+        d.bridge.go   = d.ctx->device->device.createSemaphore(sci);
+    }
+    comm->thread = std::thread(ggml_vk_comm_bridge_thread, comm);
+}
+
+static void ggml_vk_comm_bridge_free(vk_comm_device & d) {
+    if (d.bridge.done) {
+        d.ctx->device->device.destroySemaphore(d.bridge.done);
+    }
+    if (d.bridge.go) {
+        d.ctx->device->device.destroySemaphore(d.bridge.go);
+    }
+}
+
+static void ggml_vk_comm_barrier_to_host(vk_context & c) {
+    const vk::MemoryBarrier to_host{ vk::AccessFlagBits::eTransferWrite | vk::AccessFlagBits::eShaderWrite, vk::AccessFlagBits::eHostRead };
+    c->s->buffer->buf.pipelineBarrier(c->p->q->stage_flags, vk::PipelineStageFlagBits::eHost, {}, { to_host }, {}, {});
+}
+
+// the peers' slots were made available to the host before the wait resolved, the submit predates that
+static void ggml_vk_comm_barrier_from_host(vk_context & c) {
+    const vk::MemoryBarrier from_host{ vk::AccessFlagBits::eHostWrite, vk::AccessFlagBits::eShaderRead | vk::AccessFlagBits::eTransferRead };
+    c->s->buffer->buf.pipelineBarrier(vk::PipelineStageFlagBits::eHost, c->p->q->stage_flags, {}, { from_host }, {}, {});
+}
+
+static void ggml_vk_comm_bridge_signal(ggml_backend_vk_comm_context * comm, size_t i, vk_context & c, uint64_t round) {
+    const vk_comm_device & d = comm->devs[i];
+    ggml_vk_comm_barrier_to_host(c);
+    ggml_vk_comm_submit(d.ctx, c, d.bridge.done, round);
+}
+
+static void ggml_vk_comm_bridge_wait(ggml_backend_vk_comm_context * comm, size_t i, vk_context & c, uint64_t round) {
+    c->s->wait_semaphores.push_back({ comm->devs[i].bridge.go, round });
+    ggml_vk_comm_barrier_from_host(c);
+}
+
+static bool ggml_vk_comm_sfd_supported(vk_comm_device & d) {
+    const vk_device & dev = d.ctx->device;
+    const bool mesa = dev->driver_id == vk::DriverId::eMesaRadv ||
+                      dev->driver_id == vk::DriverId::eIntelOpenSourceMESA ||
+                      dev->driver_id == vk::DriverId::eMesaNvk;
+    if (!mesa || !dev->external_semaphore) {
+        return false;
+    }
+    d.sfd.get_fd    = (PFN_vkGetSemaphoreFdKHR)    dev->device.getProcAddr("vkGetSemaphoreFdKHR");
+    d.sfd.import_fd = (PFN_vkImportSemaphoreFdKHR) dev->device.getProcAddr("vkImportSemaphoreFdKHR");
+    if (!d.sfd.get_fd || !d.sfd.import_fd) {
+        return false;
+    }
+    vk::PhysicalDeviceExternalSemaphoreInfo esi{ vk::ExternalSemaphoreHandleTypeFlagBits::eSyncFd };
+    vk::ExternalSemaphoreProperties esp = dev->physical_device.getExternalSemaphoreProperties(esi);
+    const vk::ExternalSemaphoreFeatureFlags need =
+        vk::ExternalSemaphoreFeatureFlagBits::eExportable | vk::ExternalSemaphoreFeatureFlagBits::eImportable;
+    return (esp.externalSemaphoreFeatures & need) == need &&
+           (esp.compatibleHandleTypes & vk::ExternalSemaphoreHandleTypeFlagBits::eSyncFd);
+}
+
+// Mesa moves a queue to its submit thread on the first wait-before-signal, taking the CS ioctls off this thread
+static void ggml_vk_comm_sfd_enable_submit_thread(const vk_device & dev) {
+    vk::SemaphoreTypeCreateInfo stci{ vk::SemaphoreType::eTimeline, 0 };
+    vk::SemaphoreCreateInfo     sci{};
+    sci.pNext = &stci;
+    const vk::Semaphore sem   = dev->device.createSemaphore(sci);
+    const vk::Fence     fence = dev->device.createFence({});
+    const uint64_t      one   = 1;
+    vk::TimelineSemaphoreSubmitInfo tsi{ 1, &one, 0, nullptr };
+    const vk::PipelineStageFlags stage = vk::PipelineStageFlagBits::eAllCommands;
+    vk::SubmitInfo si{ 1, &sem, &stage, 0, nullptr, 0, nullptr };
+    si.setPNext(&tsi);
+    dev->compute_queue->handle->submit(si, fence);
+    dev->device.signalSemaphore({ sem, 1 });
+    VK_CHECK(dev->device.waitForFences({ fence }, true, UINT64_MAX), "ggml_vk_comm_sfd_enable_submit_thread", dev);
+    dev->device.destroyFence(fence);
+    dev->device.destroySemaphore(sem);
+}
+
+// Requires a submitted signal on sem; the export resets it.
+static int ggml_vk_comm_sfd_export(vk_comm_device & d) {
+    int fd = -1;
+    VkSemaphoreGetFdInfoKHR gi{};
+    gi.sType      = VK_STRUCTURE_TYPE_SEMAPHORE_GET_FD_INFO_KHR;
+    gi.semaphore  = d.sfd.sem;
+    gi.handleType = VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_SYNC_FD_BIT;
+    const VkResult result = d.sfd.get_fd((VkDevice) d.ctx->device->device, &gi, &fd);
+    if (result != VK_SUCCESS) {
+        GGML_ABORT("ggml_vulkan: vkGetSemaphoreFdKHR failed: %d", (int) result);
+    }
+    return fd;
+}
+
+// fd -1 imports an already signalled payload
+static void ggml_vk_comm_sfd_import(vk_comm_device & d, vk::Semaphore sem, int fd) {
+    VkImportSemaphoreFdInfoKHR isi{};
+    isi.sType      = VK_STRUCTURE_TYPE_IMPORT_SEMAPHORE_FD_INFO_KHR;
+    isi.semaphore  = sem;
+    isi.flags      = VK_SEMAPHORE_IMPORT_TEMPORARY_BIT;
+    isi.handleType = VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_SYNC_FD_BIT;
+    isi.fd         = fd;
+    const VkResult result = d.sfd.import_fd((VkDevice) d.ctx->device->device, &isi);
+    if (result != VK_SUCCESS) {
+        GGML_ABORT("ggml_vulkan: vkImportSemaphoreFdKHR failed: %d", (int) result);
+    }
+}
+
+static void ggml_vk_comm_sfd_exchange(ggml_backend_vk_comm_context * comm, uint64_t r) {
+    const size_t n = comm->devs.size();
+    for (size_t j = 0; j < n; j++) {
+        const int    fd   = ggml_vk_comm_sfd_export(comm->devs[j]);
+        const size_t last = j == n - 1 ? n - 2 : n - 1;
+        for (size_t i = 0; i < n; i++) {
+            if (i == j) {
+                continue;
+            }
+            const int imp_fd = fd < 0 || i == last ? fd : fcntl(fd, F_DUPFD_CLOEXEC, 0);
+            GGML_ASSERT(imp_fd >= 0 || fd < 0);
+            ggml_vk_comm_sfd_import(comm->devs[i], comm->devs[i].sfd.imp[j], imp_fd);
+        }
+    }
+    comm->imported.store(r, std::memory_order_release);
+}
+
+static void ggml_vk_comm_sfd_thread(ggml_backend_vk_comm_context * comm) {
+#if defined(__linux__)
+    pthread_setname_np(pthread_self(), "ggml-vk-comm");
+#endif
+    for (uint64_t r = 1; ggml_vk_comm_wait_submitted(comm, r); r++) {
+        ggml_vk_comm_sfd_exchange(comm, r);
+    }
+}
+
+static void ggml_vk_comm_sfd_init(ggml_backend_vk_comm_context * comm) {
+    vk::ExportSemaphoreCreateInfo esci{ vk::ExternalSemaphoreHandleTypeFlagBits::eSyncFd };
+    vk::SemaphoreCreateInfo       sci{};
+    sci.pNext = &esci;
+    for (vk_comm_device & d : comm->devs) {
+        d.sfd.sem = d.ctx->device->device.createSemaphore(sci);
+        d.sfd.imp.assign(comm->devs.size(), {});
+        ggml_vk_comm_sfd_enable_submit_thread(d.ctx->device);
+    }
+    comm->thread = std::thread(ggml_vk_comm_sfd_thread, comm);
+}
+
+static void ggml_vk_comm_sfd_free(vk_comm_device & d) {
+    if (d.sfd.sem) {
+        d.ctx->device->device.destroySemaphore(d.sfd.sem);
+    }
+}
+
+static void ggml_vk_comm_sfd_signal(ggml_backend_vk_comm_context * comm, size_t i, vk_context & c) {
+    vk_comm_device & d = comm->devs[i];
+    ggml_vk_comm_barrier_to_host(c);
+    ggml_vk_comm_submit(d.ctx, c, d.sfd.sem, 0);
+    for (size_t j = 0; j < comm->devs.size(); j++) {
+        d.sfd.imp[j] = j == i ? vk::Semaphore{} : ggml_vk_create_binary_semaphore(d.ctx);
+    }
+}
+
+static void ggml_vk_comm_sfd_wait(ggml_backend_vk_comm_context * comm, size_t i, vk_context & c, uint64_t round) {
+    for (size_t j = 0; j < comm->devs.size(); j++) {
+        if (j != i) {
+            c->s->wait_semaphores.push_back({ comm->devs[i].sfd.imp[j], 0 });
+        }
+    }
+    c->s->ready       = &comm->imported;
+    c->s->ready_value = round;
+    ggml_vk_comm_barrier_from_host(c);
+}
+
+void ggml_backend_vk_comm_free(void * comm_ctx) {
+    ggml_backend_vk_comm_context * comm = static_cast<ggml_backend_vk_comm_context *>(comm_ctx);
+    // pending reductions need the thread
+    for (vk_comm_device & d : comm->devs) {
+        ggml_backend_synchronize(d.backend);
+    }
+    if (comm->thread.joinable()) {
+        {
+            std::lock_guard<std::mutex> lock(comm->mutex);
+            comm->stop.store(true);
+        }
+        comm->cv.notify_all();
+        comm->thread.join();
+    }
+    for (vk_comm_device & d : comm->devs) {
+        ggml_vk_comm_sfd_free(d);
+        ggml_vk_comm_bridge_free(d);
+    }
+    ggml_vk_comm_free_buffers(comm);
+    if (comm->tctx) {
+        ggml_free(comm->tctx);
+    }
+    delete comm;
+}
+
+void * ggml_backend_vk_comm_init(ggml_backend_t * backends, size_t n_backends) {
+    if (vk_instance.allreduce_disabled || n_backends < 2) {
+        return nullptr;
+    }
+    for (size_t i = 0; i < n_backends; i++) {
+        if (!ggml_backend_is_vk(backends[i])) {
+            return nullptr;
+        }
+    }
+    ggml_backend_vk_comm_context * comm = new ggml_backend_vk_comm_context;
+    comm->devs.resize(n_backends);
+    for (size_t i = 0; i < n_backends; i++) {
+        vk_comm_device & d = comm->devs[i];
+        d.backend = backends[i];
+        d.ctx     = (ggml_backend_vk_context *) backends[i]->context;
+        d.host_buf.assign(n_backends, nullptr);
+    }
+
+    for (size_t i = 0; i < n_backends; i++) {
+        const vk_device & dev = comm->devs[i].ctx->device;
+        comm->align = std::max(comm->align, (size_t) dev->min_imported_host_pointer_alignment);
+        if (!dev->external_memory_host) {
+            GGML_LOG_WARN("ggml_vulkan: device %zu: VK_EXT_external_memory_host missing, using generic all-reduce\n", i);
+            ggml_backend_vk_comm_free(comm);
+            return nullptr;
+        }
+        if (dev->serialize_submissions) {
+            GGML_LOG_WARN("ggml_vulkan: device %zu: GGML_VK_SERIALIZE_SUBMISSIONS set, using generic all-reduce\n", i);
+            ggml_backend_vk_comm_free(comm);
+            return nullptr;
+        }
+    }
+
+    if (std::all_of(comm->devs.begin(), comm->devs.end(), ggml_vk_comm_sfd_supported)) {
+        comm->mode = vk_comm_sync::sync_fd;
+    }
+
+    if (const char * spin = getenv("GGML_VK_ALLREDUCE_SPIN_US")) {
+        comm->spin = std::chrono::microseconds(std::max(0, atoi(spin)));
+    }
+    const ggml_init_params ip = { ggml_tensor_overhead(), nullptr, true };
+    comm->tctx = ggml_init(ip);
+    comm->slot = ggml_new_tensor_1d(comm->tctx, GGML_TYPE_F32, 1);
+
+    try {
+        if (comm->mode == vk_comm_sync::sync_fd) {
+            ggml_vk_comm_sfd_init(comm);
+        } else {
+            ggml_vk_comm_bridge_init(comm);
+        }
+    } catch (const vk::SystemError & e) {
+        GGML_LOG_WARN("ggml_vulkan: semaphore initialization failed: %s, using generic all-reduce\n", e.what());
+        ggml_backend_vk_comm_free(comm);
+        return nullptr;
+    }
+    return comm;
+}
+
+// A: partial to own host slot (round & 1), signal; B: wait on peers' A, add their slots; all A precede any B
+// slot reuse: A(k) follows own B(k-1), which waited on peers' A(k-1), submitted after their B(k-2), the last reader
+bool ggml_backend_vk_comm_allreduce_tensor(void * comm_ctx, ggml_tensor ** tensors) {
+    ggml_backend_vk_comm_context * comm = static_cast<ggml_backend_vk_comm_context *>(comm_ctx);
+    const size_t n = comm->devs.size();
+
+    const int64_t ne = ggml_nelements(tensors[0]);
+    for (size_t i = 0; i < n; i++) {
+        const ggml_tensor * t = tensors[i];
+        if (t->type != GGML_TYPE_F32 || ggml_nelements(t) != ne || !ggml_is_contiguous(t) || !ggml_backend_buffer_is_vk(t->buffer)) {
+            return false;
+        }
+        const ggml_backend_vk_context * ctx = comm->devs[i].ctx;
+        if (!ctx->compute_ctx.expired() && !ctx->compute_ctx.lock()->out_memcpys.empty()) {
+            return false;
+        }
+    }
+    if (ne == 0) {
+        return true;
+    }
+    const size_t nbytes = ggml_nbytes(tensors[0]);
+    if (!ggml_backend_vk_comm_ensure(comm, nbytes)) {
+        return false;
+    }
+
+    const bool      f32_wire  = nbytes < ggml_vk_comm_f16_min_size || ggml_vk_comm_needs_f32(tensors[0]);
+    const ggml_type wire      = f32_wire ? GGML_TYPE_F32 : GGML_TYPE_F16;
+    const uint64_t  round     = ++comm->round;
+    const size_t    slot_offs = (round & 1) * comm->cap;
+    const size_t    slot_size = ne * ggml_type_size(wire);
+
+    for (size_t i = 0; i < n; i++) {
+        const vk_comm_device &    d   = comm->devs[i];
+        ggml_backend_vk_context * ctx = d.ctx;
+        ggml_tensor *             t   = tensors[i];
+        vk_buffer & tbuf = ((ggml_backend_vk_buffer_context *) t->buffer->context)->dev_buffer;
+        vk_buffer & hbuf = ((ggml_backend_vk_buffer_context *) d.host_buf[i]->context)->dev_buffer;
+        const size_t toffs = vk_tensor_offset(t) + t->view_offs;
+
+        vk_context c = ggml_vk_get_compute_ctx(ctx);
+        ggml_vk_sync_buffers(ctx, c);
+        if (!(t->flags & GGML_TENSOR_FLAG_COMPUTE)) {
+            c->s->buffer->buf.fillBuffer(hbuf->buffer, slot_offs, GGML_PAD(slot_size, 4), 0);
+            c->s->buffer->buf.fillBuffer(tbuf->buffer, toffs, nbytes, 0);
+        } else if (wire == GGML_TYPE_F32) {
+            ggml_vk_buffer_copy_async(c, hbuf, slot_offs, tbuf, toffs, nbytes);
+        } else {
+            ggml_vk_comm_set_slot(comm->slot, t, wire, d.host_buf[i], slot_offs);
+            ggml_vk_cpy(ctx, c, t, comm->slot);
+        }
+        if (comm->mode == vk_comm_sync::sync_fd) {
+            ggml_vk_comm_sfd_signal(comm, i, c);
+        } else {
+            ggml_vk_comm_bridge_signal(comm, i, c, round);
+        }
+    }
+    ggml_vk_comm_notify_submitted(comm, round);
+
+    for (size_t i = 0; i < n; i++) {
+        ggml_backend_vk_context * ctx = comm->devs[i].ctx;
+        vk_context c = ggml_vk_get_compute_ctx(ctx);
+        ctx->do_add_rms_partials = false;
+        if (comm->mode == vk_comm_sync::sync_fd) {
+            ggml_vk_comm_sfd_wait(comm, i, c, round);
+        } else {
+            ggml_vk_comm_bridge_wait(comm, i, c, round);
+        }
+        ggml_vk_sync_buffers(ctx, c);
+        for (size_t j = 0; j < n; j++) {
+            if (j == i) {
+                continue;
+            }
+            ggml_vk_comm_set_slot(comm->slot, tensors[i], wire, comm->devs[j].host_buf[i], slot_offs);
+            ggml_vk_add(ctx, c, tensors[i], comm->slot, tensors[i]);
+            ggml_vk_sync_buffers(ctx, c);
+        }
+        // small reductions go out with the next graph's first submit
+        if (wire != GGML_TYPE_F32 || vk_perf_logger_enabled) {
+            ggml_vk_comm_submit(ctx, c, {}, 0);
+        }
+    }
+    return true;
+}
+#endif

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

@@ -86,6 +86,7 @@ void ggml_vk_cmd_label_begin(vk::CommandBuffer buf, const char * name);
 void ggml_vk_ctx_end(vk_context& ctx);
 void ggml_vk_ctx_begin(vk_device& device, vk_context& subctx);
 vk_context ggml_vk_get_compute_ctx(ggml_backend_vk_context * ctx);
+vk::Semaphore ggml_vk_create_binary_semaphore(ggml_backend_vk_context * ctx);
 vk_context ggml_vk_get_transfer_ctx(ggml_backend_vk_context * ctx);
 bool ggml_vk_submit_transfer_ctx(ggml_backend_vk_context * ctx);
 size_t ggml_vk_align_size(size_t width, size_t align);
@@ -233,8 +234,16 @@ size_t ggml_backend_vk_buffer_type_get_alignment(ggml_backend_buffer_type_t buft
 size_t ggml_backend_vk_buffer_type_get_max_size(ggml_backend_buffer_type_t buft);
 size_t ggml_backend_vk_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor);
 void ggml_backend_vk_free(ggml_backend_t backend);
+ggml_backend_buffer_t ggml_backend_vk_device_buffer_from_host_ptr(ggml_backend_dev_t dev, void * ptr, size_t size, size_t max_tensor_size);
 ggml_backend_reg_t ggml_backend_vk_reg();
 
+// comm
+#if !defined(_WIN32)
+void * ggml_backend_vk_comm_init(ggml_backend_t * backends, size_t n_backends);
+void ggml_backend_vk_comm_free(void * comm_ctx);
+bool ggml_backend_vk_comm_allreduce_tensor(void * comm_ctx, ggml_tensor ** tensors);
+#endif
+
 // debug
 int64_t ggml_vk_get_op_batch_size(const ggml_tensor * op);
 #ifdef GGML_VULKAN_RUN_TESTS

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

@@ -679,6 +679,7 @@ struct vk_device_struct {
     uint64_t suballocation_block_size;
     uint64_t min_imported_host_pointer_alignment;
     bool external_memory_host {};
+    bool external_semaphore {};
     bool fp16;
     bool bf16;
     bool pipeline_robustness;
@@ -1087,6 +1088,9 @@ struct vk_submission {
     vk_command_buffer* buffer = nullptr;
     std::vector<vk_semaphore> wait_semaphores;
     std::vector<vk_semaphore> signal_semaphores;
+    // submit waits until *ready >= ready_value
+    const std::atomic<uint64_t> * ready = nullptr;
+    uint64_t ready_value = 0;
 };
 
 typedef std::vector<vk_submission> vk_sequence;
@@ -1225,7 +1229,7 @@ struct ggml_backend_vk_context {
 
     vk_device device;
 
-    size_t semaphore_idx, event_idx;
+    size_t semaphore_idx, binary_semaphore_idx, event_idx;
     ggml_vk_garbage_collector gc;
     size_t prealloc_size_x, prealloc_size_y, prealloc_size_split_k, prealloc_size_add_rms_partials, prealloc_size_add_rms_partials_offset;
     vk_buffer prealloc_x, prealloc_y, prealloc_split_k, prealloc_add_rms_partials, sync_staging;
@@ -1341,6 +1345,7 @@ struct vk_instance_t {
     uint32_t host_import_failed = 0;  // bitmask by device->idx, guarded by host_allocations_mutex
 
     bool multi_device = false;
+    bool allreduce_disabled = false;
 };
 
 typedef void (*ggml_vk_func_t)(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst);

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

@@ -934,6 +934,11 @@ void ggml_vk_submit(vk_context& ctx, vk::Fence fence) {
 
     for (const auto& sequence : ctx->seqs) {
         for (const auto& submission : sequence) {
+            if (submission.ready) {
+                while (submission.ready->load(std::memory_order_acquire) < submission.ready_value) {
+                    YIELD();
+                }
+            }
             stage_flags.push_back({});
             idx++;
             tl_wait_vals.push_back({});
@@ -1082,14 +1087,16 @@ vk_context ggml_vk_create_temporary_context(vk_command_pool& p) {
     return result;
 }
 
-static vk_semaphore * ggml_vk_create_binary_semaphore(ggml_backend_vk_context * ctx) {
-    VK_LOG_DEBUG("ggml_vk_create_timeline_semaphore()");
-    vk::SemaphoreTypeCreateInfo tci{ vk::SemaphoreType::eBinary, 0 };
-    vk::SemaphoreCreateInfo ci{};
-    ci.setPNext(&tci);
-    vk::Semaphore semaphore = ctx->device->device.createSemaphore(ci);
-    ctx->gc.semaphores.push_back({ semaphore, 0 });
-    return &ctx->gc.semaphores[ctx->gc.semaphores.size() - 1];
+vk::Semaphore ggml_vk_create_binary_semaphore(ggml_backend_vk_context * ctx) {
+    VK_LOG_DEBUG("ggml_vk_create_binary_semaphore()");
+    if (ctx->binary_semaphore_idx >= ctx->gc.semaphores.size()) {
+        vk::SemaphoreTypeCreateInfo tci{ vk::SemaphoreType::eBinary, 0 };
+        vk::SemaphoreCreateInfo ci{};
+        ci.setPNext(&tci);
+        vk::Semaphore semaphore = ctx->device->device.createSemaphore(ci);
+        ctx->gc.semaphores.push_back({ semaphore, 0 });
+    }
+    return ctx->gc.semaphores[ctx->binary_semaphore_idx++].s;
 }
 
 static vk_semaphore * ggml_vk_create_timeline_semaphore(ggml_backend_vk_context * ctx) {
@@ -1450,8 +1457,8 @@ static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vec
         lut_size = 4*16;
         break;
     case GGML_TYPE_NVFP4:
-        // Same kvalues budget as MXFP4 plus ue4m3_fp32_lut[128] (types.glsl, DATA_A_NVFP4).
-        lut_size = 4*16 + 128u * (uint32_t)sizeof(float);
+        // Same kvalues budget as MXFP4 plus ue4m3_fp32_lut[256] (types.glsl, DATA_A_NVFP4).
+        lut_size = 4*16 + 256u * (uint32_t)sizeof(float);
         break;
     default:
         break;
@@ -1585,7 +1592,7 @@ static bool ggml_vk_matmul_cm1_int_shmem_support(const vk_device& device, const
         total += 16u * (uint32_t)sizeof(int8_t);         // cm1_kvalues[16]
     }
     if (src0_type == GGML_TYPE_NVFP4 && !device->ocp_fp4) {
-        total += 128u * (uint32_t)sizeof(float);         // ue4m3_fp32_lut[128]
+        total += 256u * (uint32_t)sizeof(float);         // ue4m3_fp32_lut[256]
     }
     if (mul_mat_id) {
         total += BN * 2u * (uint32_t)sizeof(uint16_t);   // row_ids[BN] (u16vec2)
@@ -4128,6 +4135,8 @@ vk_device ggml_vk_get_device(size_t idx) {
                 device->memory_priority = true;
             } else if (strcmp("VK_EXT_external_memory_host", properties.extensionName) == 0) {
                 device->external_memory_host = true;
+            } else if (strcmp("VK_KHR_external_semaphore_fd", properties.extensionName) == 0) {
+                device->external_semaphore = true;
 #if defined(VK_EXT_shader_64bit_indexing)
             } else if (strcmp("VK_EXT_shader_64bit_indexing", properties.extensionName) == 0) {
                 device->shader_64b_indexing = true;
@@ -4475,6 +4484,9 @@ vk_device ggml_vk_get_device(size_t idx) {
         if (device->external_memory_host) {
             device_extensions.push_back("VK_EXT_external_memory_host");
         }
+        if (device->external_semaphore) {
+            device_extensions.push_back("VK_KHR_external_semaphore_fd");
+        }
 
 #if defined(VK_EXT_shader_64bit_indexing)
         VkPhysicalDeviceShader64BitIndexingFeaturesEXT shader_64bit_indexing_features {};
@@ -5287,6 +5299,7 @@ void ggml_vk_instance_init() {
     vk_perf_logger_concurrent = getenv("GGML_VK_PERF_LOGGER_CONCURRENT") != nullptr;
     vk_enable_sync_logger = getenv("GGML_VK_SYNC_LOGGER") != nullptr;
     vk_memory_logger_enabled = getenv("GGML_VK_MEMORY_LOGGER") != nullptr;
+    vk_instance.allreduce_disabled = getenv("GGML_VK_DISABLE_ALLREDUCE") != nullptr;
     const char* GGML_VK_PIPELINE_STATS = getenv("GGML_VK_PIPELINE_STATS");
     if (GGML_VK_PIPELINE_STATS != nullptr) {
         vk_pipeline_stats_filter = GGML_VK_PIPELINE_STATS;
@@ -5493,6 +5506,7 @@ void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) {
     ctx->device = ggml_vk_get_device(idx);
 
     ctx->semaphore_idx = 0;
+    ctx->binary_semaphore_idx = 0;
     ctx->event_idx = 0;
 
     ctx->prealloc_size_x = 0;
@@ -5920,8 +5934,18 @@ static uint32_t ggml_vk_guess_split_k(ggml_backend_vk_context * ctx, uint32_t m,
         return 1;
     }
 
+    const bool rdna = ctx->device->architecture == AMD_RDNA3 || ctx->device->architecture == AMD_RDNA4;
     uint32_t split_k = 1;
-    if (ctx->device->shader_core_count != 0 && n >= pipeline->wg_denoms[1]) {
+    if (rdna && ctx->device->shader_core_count != 0 && k >= 2048 && n <= 32) {
+        // Small n gives too few workgroups to hide latency, aim for 8 per CU. Small tiles are one subgroup each.
+        const uint32_t tiles = CEIL_DIV(m, pipeline->wg_denoms[0]) * CEIL_DIV(n, pipeline->wg_denoms[1]);
+        if (tiles < 4 * ctx->device->shader_core_count) {
+            split_k = std::min(8u, CEIL_DIV(8 * ctx->device->shader_core_count, tiles));
+        }
+    }
+
+    // On RDNA, also split when n <= 64 only partly fills one tile column.
+    if (split_k == 1 && ctx->device->shader_core_count != 0 && (n >= pipeline->wg_denoms[1] || (rdna && n <= 64))) {
         // If k is 'large' and the SMs will fill less than halfway, use split_k.
         uint32_t m_tiles = CEIL_DIV(m, pipeline->wg_denoms[0]);
         uint32_t n_tiles = CEIL_DIV(n, pipeline->wg_denoms[1]);
@@ -5934,24 +5958,24 @@ static uint32_t ggml_vk_guess_split_k(ggml_backend_vk_context * ctx, uint32_t m,
             }
             // Cap the split at 8x. Unless k is huge this is a lot of overhead.
             split_k = std::min(split_k, 8u);
-
-            // ggml_vk_matmul will align the splits to be a multiple of 256.
-            // If this rounded up size would cause the last split to be empty,
-            // then reduce the split count.
-            while (true) {
-                if (split_k == 1) {
-                    break;
-                }
-                uint32_t k_split = CEIL_DIV(k, split_k);
-                k_split = ROUNDUP_POW2(k_split, 256);
-                if (k_split * (split_k - 1) < k) {
-                    break;
-                }
-                split_k--;
-            }
         }
     }
 
+    // ggml_vk_matmul will align the splits to be a multiple of 256.
+    // If this rounded up size would cause the last split to be empty,
+    // then reduce the split count.
+    while (true) {
+        if (split_k == 1) {
+            break;
+        }
+        uint32_t k_split = CEIL_DIV(k, split_k);
+        k_split = ROUNDUP_POW2(k_split, 256);
+        if (k_split * (split_k - 1) < k) {
+            break;
+        }
+        split_k--;
+    }
+
     return split_k;
 }
 
@@ -13003,10 +13027,7 @@ void ggml_vk_graph_cleanup(ggml_backend_vk_context * ctx) {
         ggml_vk_command_pool_cleanup(ctx->device, ctx->transfer_cmd_pool);
     }
 
-    for (size_t i = 0; i < ctx->gc.semaphores.size(); i++) {
-        ctx->device->device.destroySemaphore({ ctx->gc.semaphores[i].s });
-    }
-    ctx->gc.semaphores.clear();
+    ctx->binary_semaphore_idx = 0;
 
     for (size_t i = 0; i < ctx->gc.tl_semaphores.size(); i++) {
         ctx->device->device.destroySemaphore({ ctx->gc.tl_semaphores[i].s });
@@ -13054,6 +13075,11 @@ void ggml_vk_cleanup(ggml_backend_vk_context * ctx) {
     }
     ctx->gc.events.clear();
 
+    for (auto & sem : ctx->gc.semaphores) {
+        ctx->device->device.destroySemaphore(sem.s);
+    }
+    ctx->gc.semaphores.clear();
+
     ctx->device->device.destroyFence(ctx->fence);
     ctx->device->device.destroyFence(ctx->almost_ready_fence);
 
@@ -14454,7 +14480,9 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
             flops_cap = 2'000'000'000ULL * ctx->device->shader_core_count;
         }
     }
-    uint64_t flops_per_submit = std::min(flops_cap, ctx->last_total_flops / 40u);
+    // small graphs gain nothing from splitting, each submit costs CPU time
+    uint64_t flops_per_submit = ctx->last_total_flops < 1'000'000'000ULL ? 0 :
+        std::clamp(ctx->last_total_flops / 40u, std::min<uint64_t>(100'000'000ULL, flops_cap), flops_cap);
 
     auto const submit_after = [&](int start, int end) {
         if (ctx->device->serialize_submissions) {
@@ -15522,13 +15550,14 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
         }
         return true;
     };
-    // reject any tensors larger than the max buffer size
+    // non-compute ops are never executed here; their size is bounded by the buffer allocation limit
+    const bool non_compute = ggml_op_is_empty(op->op);
     for (int i = 0; i < GGML_MAX_SRC; i++) {
-        if (op->src[i] && !tensor_size_supported(ggml_nbytes(op->src[i]))) {
+        if (op->src[i] && !non_compute && !tensor_size_supported(ggml_nbytes(op->src[i]))) {
             return false;
         }
     }
-    if (!tensor_size_supported(ggml_nbytes(op))) {
+    if (!non_compute && !tensor_size_supported(ggml_nbytes(op))) {
         return false;
     }
 
@@ -15654,6 +15683,10 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
             }
         case GGML_OP_FLASH_ATTN_EXT:
             {
+                // K/V row lists (ggml_flash_attn_ext_rows) are not supported
+                if (op->src[5] != nullptr) {
+                    return false;
+                }
                 bool coopmat2 = device->coopmat2;
                 uint32_t HSK = op->src[1]->ne[0];
                 uint32_t HSV = op->src[2]->ne[0];
@@ -16256,7 +16289,7 @@ static void ggml_backend_vk_device_event_synchronize(ggml_backend_dev_t dev, ggm
     }
 }
 
-static ggml_backend_buffer_t ggml_backend_vk_device_buffer_from_host_ptr(ggml_backend_dev_t dev, void * ptr, size_t size, size_t max_tensor_size) {
+ggml_backend_buffer_t ggml_backend_vk_device_buffer_from_host_ptr(ggml_backend_dev_t dev, void * ptr, size_t size, size_t max_tensor_size) {
     VK_LOG_DEBUG("ggml_backend_vk_device_buffer_from_host_ptr(backend=" << dev << ", ptr=" << ptr << ", size=" << size << ")");
     GGML_UNUSED(max_tensor_size);
 
@@ -16338,11 +16371,29 @@ static ggml_backend_dev_t ggml_backend_vk_reg_get_device(ggml_backend_reg_t reg,
     return devices[device];
 }
 
+static void * ggml_backend_vk_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) {
+    UNUSED(reg);
+#if !defined(_WIN32)
+    if (strcmp(name, "ggml_backend_comm_init") == 0) {
+        return (void *) ggml_backend_vk_comm_init;
+    }
+    if (strcmp(name, "ggml_backend_comm_free") == 0) {
+        return (void *) ggml_backend_vk_comm_free;
+    }
+    if (strcmp(name, "ggml_backend_comm_allreduce_tensor") == 0) {
+        return (void *) ggml_backend_vk_comm_allreduce_tensor;
+    }
+#else
+    UNUSED(name);
+#endif
+    return nullptr;
+}
+
 static const struct ggml_backend_reg_i ggml_backend_vk_reg_i = {
     /* .get_name         = */ ggml_backend_vk_reg_get_name,
     /* .get_device_count = */ ggml_backend_vk_reg_get_device_count,
     /* .get_device       = */ ggml_backend_vk_reg_get_device,
-    /* .get_proc_address = */ NULL,
+    /* .get_proc_address = */ ggml_backend_vk_reg_get_proc_address,
 };
 
 ggml_backend_reg_t ggml_backend_vk_reg() {

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

@@ -40,6 +40,8 @@ layout (binding = 2) writeonly buffer DST {float data_dst[];};
 shared uint vals[EXPERTS_PER_CHUNK];
 shared uint offsets[EXPERTS_PER_CHUNK];
 shared uint cursors[EXPERTS_PER_CHUNK];
+shared uint skipped_rows;
+shared uint skip_count;
 
 // data_d[0] is the tile count. Each tile has (expert, valid rows).
 // Packed (token << 16) | slot row ids start at p.row_ids_offset, with each expert on a tile boundary.
@@ -51,23 +53,38 @@ void write_tiles(uint expert, uint count, uint tile_begin, uint n_tiles) {
     }
 }
 
-// the matmul writes only the rows that an expert owns, so zero the rows of the skipped slots
+// The matmul does not write skipped rows, so collect their offsets and zero them cooperatively.
 void zero_skipped_rows() {
     const uint num_elements = p.ne00 * p.ne01;
     const uint tid = gl_LocalInvocationID.x;
 
-    for (uint idx = 0; idx < num_elements; idx++) {
-        const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L);
-        const uint i00 = idx - i01 * p.ne00;
-
-        if (int(data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00]) != -1) {
-            continue;
+    // Reuse offsets[] as a queue, so each batch must fit in EXPERTS_PER_CHUNK.
+    for (uint base = 0; base < num_elements; base += EXPERTS_PER_CHUNK) {
+        const uint end = min(base + EXPERTS_PER_CHUNK, num_elements);
+        if (tid == 0) {
+            skip_count = 0;
         }
+        barrier();
 
-        const uint d_off = i00 * p.stride_dst + i01 * p.batch_stride_dst;
-        for (uint i = tid; i < p.ne0_dst; i += BLOCK_SIZE) {
-            data_dst[d_off + i] = 0.0f;
+        for (uint idx = base + tid; idx < end; idx += BLOCK_SIZE) {
+            const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L);
+            const uint i00 = idx - i01 * p.ne00;
+
+            if (int(data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00]) == -1) {
+                const uint slot = atomicAdd(skip_count, 1);
+                offsets[slot] = i00 * p.stride_dst + i01 * p.batch_stride_dst;
+            }
         }
+        barrier();
+
+        const uint n_skipped = skip_count;
+        for (uint r = 0; r < n_skipped; ++r) {
+            const uint d_off = offsets[r];
+            for (uint i = tid; i < p.ne0_dst; i += BLOCK_SIZE) {
+                data_dst[d_off + i] = 0.0f;
+            }
+        }
+        barrier();
     }
 }
 
@@ -76,6 +93,11 @@ void main() {
     const uint tid = gl_LocalInvocationID.x;
     const uint tile_shift = findLSB(p.row_tile_size);
     uint total_tiles = 0;
+    uint local_skipped = 0;
+    if (tid == 0) {
+        skipped_rows = 0;
+    }
+    barrier();
 #ifdef USE_SUBGROUPS
     // Use the subgroup that contains invocation 0.
     const bool prefix_subgroup = subgroupAny(tid == 0);
@@ -135,7 +157,12 @@ void main() {
         for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) {
             const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L);
             const uint i00 = idx - i01 * p.ne00;
-            const uint expert = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00] - expert_base;
+            const uint id = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00];
+            // Count skipped slots only in the first expert chunk.
+            if (expert_base == 0 && int(id) == -1) {
+                local_skipped++;
+            }
+            const uint expert = id - expert_base;
             if (expert < n_experts) {
                 const uint row = atomicAdd(cursors[expert], 1);
                 data_d[p.row_ids_offset + offsets[expert] + row] = (i01 << 16) | i00;
@@ -144,9 +171,17 @@ void main() {
         // The next chunk clears vals. Scatter uses only offsets and cursors.
         expert_base += n_experts;
     }
+    if (local_skipped != 0) {
+        atomicAdd(skipped_rows, local_skipped);
+    }
+    // All scatter reads must finish before reusing offsets[] as a queue.
+    barrier();
+
     if (tid == 0) {
         data_d[0] = total_tiles;
     }
 
-    zero_skipped_rows();
+    if (skipped_rows != 0) {
+        zero_skipped_rows();
+    }
 }

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

@@ -173,7 +173,7 @@ void main() {
         cm1_kvalues[gl_LocalInvocationIndex] = kvalues_mxfp4_const[gl_LocalInvocationIndex];
     }
 #if !defined(USE_OCP_FP4)
-    for (uint i = gl_LocalInvocationIndex; i < 128u; i += BLOCK_SIZE) {
+    for (uint i = gl_LocalInvocationIndex; i < 256u; i += BLOCK_SIZE) {
         ue4m3_fp32_lut[i] = ue4m3_to_fp32_build(i);
     }
 #endif

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

@@ -1921,8 +1921,9 @@ shared int8_t kvalues_mxfp4[16];
 #endif
 
 #if defined(DATA_A_NVFP4) && !defined(USE_OCP_FP4)
-// UE4M3 scale in NVFP4 blocks use only 7 bits; sign (bit 7) is always zero.
-shared float ue4m3_fp32_lut[128];
+// Cover all scale bytes accepted by ggml_validate_row_data.
+// ue4m3_to_fp32_build() masks the exponent and mantissa, but checks the full byte for 0 and 0x7F.
+shared float ue4m3_fp32_lut[256];
 
 float ue4m3_to_fp32_build(uint u) {
     if (u == 0u || u == 127u) {
@@ -1947,7 +1948,7 @@ void init_iq_shmem(uvec3 wgsize)
         kvalues_mxfp4[i] = kvalues_mxfp4_const[i];
     }
 #if defined(DATA_A_NVFP4)
-    for (uint i = gl_LocalInvocationIndex.x; i < 128u; i += wgsize.x) {
+    for (uint i = gl_LocalInvocationIndex.x; i < 256u; i += wgsize.x) {
         ue4m3_fp32_lut[i] = ue4m3_to_fp32_build(i);
     }
 #endif

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

@@ -19,6 +19,7 @@
 #define GGML_WEBGPU_I32_SIZE_BYTES                   4
 #define GGML_WEBGPU_FLASH_ATTN_PREFERRED_KV_SG_TILES 8u
 #define GGML_WEBGPU_FLASH_ATTN_VEC_MAX_SEQ_LEN       20u
+#define GGML_WEBGPU_MUL_MAT_VEC_MAX_ROWS             4
 #define GGML_WEBGPU_FLASH_ATTN_VEC_MAX_KV_TILE       32u
 #define GGML_WEBGPU_FLASH_ATTN_TILE_MAX_KV_TILE      64u
 #define GGML_WEBGPU_FLASH_ATTN_PREFERRED_WG_SIZE     128u
@@ -1189,9 +1190,10 @@ struct ggml_webgpu_soft_max_pipeline_key_hash {
 
 inline bool ggml_webgpu_can_use_mmvq(const ggml_tensor * src0,
                                      const ggml_tensor * src1,
+                                     int64_t             n_rows,
                                      bool                supports_dot_product,
                                      const std::string & vendor) {
-    if (src1->ne[1] <= 4) {
+    if (n_rows <= GGML_WEBGPU_MUL_MAT_VEC_MAX_ROWS) {
         bool supports_dp4a = vendor == "amd" || vendor == "intel" || vendor == "nvidia";
         if (supports_dp4a && supports_dot_product) {
             switch (src1->type) {
@@ -1231,6 +1233,7 @@ class ggml_webgpu_shader_lib {
     std::unordered_map<int, webgpu_pipeline> argsort_pipelines;        // key is order
     std::unordered_map<int, webgpu_pipeline> argsort_merge_pipelines;  // key is order
     std::unordered_map<int, webgpu_pipeline> cumsum_pipelines;         // key is fixed, no variants yet
+    std::unordered_map<int, webgpu_pipeline> arange_pipelines;         // key is fixed, no variants yet
     std::unordered_map<ggml_webgpu_row_norm_pipeline_key, webgpu_pipeline, ggml_webgpu_row_norm_pipeline_key_hash>
         row_norm_pipelines;                                            // op/inplace
 
@@ -1941,6 +1944,24 @@ class ggml_webgpu_shader_lib {
         return gated_delta_net_pipelines[key];
     }
 
+    webgpu_pipeline get_arange_pipeline(const ggml_webgpu_shader_lib_context & context) {
+        auto it = arange_pipelines.find(1);
+        if (it != arange_pipelines.end()) {
+            return it->second;
+        }
+
+        std::vector<std::string> defines;
+        defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size));
+
+        auto processed           = preprocessor.preprocess(wgsl_arange, defines);
+        auto decisions           = std::make_shared<ggml_webgpu_generic_shader_decisions>();
+        decisions->wg_size       = context.max_wg_size;
+        webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, "arange");
+        pipeline.context         = decisions;
+        arange_pipelines[1]      = pipeline;
+        return arange_pipelines[1];
+    }
+
     webgpu_pipeline get_pad_pipeline(const ggml_webgpu_shader_lib_context & context) {
         ggml_webgpu_pad_pipeline_key key = {};
         key.circular                     = ggml_get_op_params_i32(context.dst, 8) != 0;
@@ -2011,13 +2032,13 @@ class ggml_webgpu_shader_lib {
         ggml_webgpu_mul_mat_vec_pipeline_key key = {};
         key.src0_type                            = context.src0->type;
         key.src1_type                            = context.src1->type;
-        key.vectorized = (context.src0->ne[0] % 4 == 0 &&
-                          (context.src0->type == GGML_TYPE_F32 || context.src0->type == GGML_TYPE_F16)) ?
-                             1 :
-                             0;
-        key.num_cols   = context.dst->ne[1];
-        key.use_mmvq =
-            ggml_webgpu_can_use_mmvq(context.src0, context.src1, context.supports_dot_product, context.vendor);
+        key.vectorized  = (context.src0->ne[0] % 4 == 0 &&
+                           (context.src0->type == GGML_TYPE_F32 || context.src0->type == GGML_TYPE_F16)) ?
+                              1 :
+                              0;
+        key.num_cols    = context.dst->ne[1];
+        key.use_mmvq    = ggml_webgpu_can_use_mmvq(context.src0, context.src1, context.src1->ne[1],
+                                                   context.supports_dot_product, context.vendor);
         key.src_overlap = src_overlap;
 
         auto it = mul_mat_vec_pipelines.find(key);

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

@@ -856,6 +856,32 @@ static webgpu_encoded_op ggml_webgpu_set(webgpu_context & ctx,
     return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y);
 }
 
+static webgpu_encoded_op ggml_webgpu_arange(webgpu_context & ctx, ggml_tensor * dst) {
+    ggml_webgpu_shader_lib_context shader_lib_ctx = {};
+    shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup;
+
+    webgpu_pipeline pipeline  = ctx->shader_lib->get_arange_pipeline(shader_lib_ctx);
+    auto *          decisions = static_cast<ggml_webgpu_generic_shader_decisions *>(pipeline.context.get());
+
+    const uint32_t        ne     = (uint32_t) ggml_nelements(dst);
+    std::vector<uint32_t> params = {
+        ne,
+        (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)),
+        ggml_webgpu_u32_from_f32(ggml_get_op_params_f32(dst, 0)),
+        ggml_webgpu_u32_from_f32(ggml_get_op_params_f32(dst, 2)),
+    };
+
+    std::vector<wgpu::BindGroupEntry> entries = {
+        ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, dst),
+    };
+
+    uint32_t wg_x;
+    uint32_t wg_y;
+    uint32_t total_wg = CEIL_DIV(ne, decisions->wg_size);
+    compute_2d_workgroups(total_wg, ctx->global_ctx->capabilities.limits.maxComputeWorkgroupsPerDimension, wg_x, wg_y);
+    return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y);
+}
+
 static webgpu_encoded_op ggml_webgpu_pad(webgpu_context & ctx, ggml_tensor * src, ggml_tensor * dst) {
     ggml_webgpu_shader_lib_context shader_lib_ctx = {};
     shader_lib_ctx.src0                           = src;
@@ -1634,11 +1660,11 @@ static webgpu_encoded_op ggml_webgpu_mul_mat(webgpu_context & ctx,
                                              ggml_tensor *    src1,
                                              ggml_tensor *    dst) {
     // Determine if this is a mat-vec operation
-    bool use_mat_vec = (dst->ne[1] <= 4);
+    bool use_mat_vec = (dst->ne[1] <= GGML_WEBGPU_MUL_MAT_VEC_MAX_ROWS);
 
     // use MMVQ path for mat-vec
-    bool use_mmvq = ggml_webgpu_can_use_mmvq(src0, src1, ctx->global_ctx->capabilities.supports_dot_product,
-                                             ctx->global_ctx->vendor);
+    bool use_mmvq = ggml_webgpu_can_use_mmvq(
+        src0, src1, src1->ne[1], ctx->global_ctx->capabilities.supports_dot_product, ctx->global_ctx->vendor);
 
     ggml_webgpu_shader_lib_context shader_lib_ctx = {};
 
@@ -1965,7 +1991,8 @@ struct ggml_webgpu_flash_attn_op {
 static bool ggml_webgpu_flash_attn_use_vec_path(const webgpu_global_context & global_ctx,
                                                 const ggml_tensor *           Q,
                                                 const ggml_tensor *           K,
-                                                const ggml_tensor *           V) {
+                                                const ggml_tensor *           V,
+                                                int64_t                       n_q) {
     const size_t storage_offset_alignment = global_ctx->capabilities.limits.minStorageBufferOffsetAlignment;
 
     const bool k_float_vec4_aligned = (K->type != GGML_TYPE_F16 && K->type != GGML_TYPE_F32) ||
@@ -1979,7 +2006,7 @@ static bool ggml_webgpu_flash_attn_use_vec_path(const webgpu_global_context & gl
         ggml_is_quantized(V->type) ? ggml_blck_size(V->type) : GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH;
     const bool kv_vec_head_dims_aligned = Q->ne[0] % k_vec_head_align == 0 && V->ne[0] % v_vec_head_align == 0;
 
-    return global_ctx->capabilities.supports_subgroups && (Q->ne[1] < GGML_WEBGPU_FLASH_ATTN_VEC_MAX_SEQ_LEN) &&
+    return global_ctx->capabilities.supports_subgroups && (n_q < GGML_WEBGPU_FLASH_ATTN_VEC_MAX_SEQ_LEN) &&
            kv_vec_head_dims_aligned && k_float_vec4_aligned && v_float_vec4_aligned;
 }
 
@@ -2288,7 +2315,7 @@ static webgpu_encoded_op ggml_webgpu_flash_attn(webgpu_context & ctx,
                                                 ggml_tensor *    sinks,
                                                 ggml_tensor *    dst) {
     ggml_webgpu_flash_attn_op op = ggml_webgpu_flash_attn_prepare(ctx, Q, K, V, mask, sinks, dst);
-    if (ggml_webgpu_flash_attn_use_vec_path(ctx->global_ctx, Q, K, V)) {
+    if (ggml_webgpu_flash_attn_use_vec_path(ctx->global_ctx, Q, K, V, Q->ne[1])) {
         return ggml_webgpu_flash_attn_vec(ctx, Q, K, V, mask, sinks, dst, std::move(op));
     }
     return ggml_webgpu_flash_attn_direct(ctx, op);
@@ -3479,6 +3506,8 @@ static std::optional<webgpu_encoded_op> ggml_webgpu_encode(webgpu_context ctx,
                 return ggml_webgpu_gated_delta_net(ctx, src0, src1, src2, node->src[3], node->src[4], node->src[5],
                                                    node, nullptr);
             }
+        case GGML_OP_ARANGE:
+            return ggml_webgpu_arange(ctx, node);
         case GGML_OP_PAD:
             return ggml_webgpu_pad(ctx, src0, node);
         case GGML_OP_ARGMAX:
@@ -4012,7 +4041,10 @@ static size_t ggml_backend_webgpu_buffer_type_get_alloc_size(ggml_backend_buffer
                 const ggml_tensor * V            = tensor->src[2];
                 const ggml_tensor * mask         = tensor->src[3];
                 const auto &        capabilities = ctx->webgpu_global_ctx->capabilities;
-                if (ggml_webgpu_flash_attn_use_vec_path(ctx->webgpu_global_ctx, Q, K, V)) {
+                // the vec path scratch grows with the query rows, so a longer batch reserves the scratch of the
+                // longest vec batch and the size stays monotonic in the batch length
+                const int64_t       n_q = std::min<int64_t>(Q->ne[1], GGML_WEBGPU_FLASH_ATTN_VEC_MAX_SEQ_LEN - 1);
+                if (ggml_webgpu_flash_attn_use_vec_path(ctx->webgpu_global_ctx, Q, K, V, n_q)) {
                     const bool kv_direct =
                         ggml_webgpu_flash_attn_k_direct(Q, K, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH) ||
                         ggml_webgpu_flash_attn_v_direct(Q, V, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH);
@@ -4024,7 +4056,7 @@ static size_t ggml_backend_webgpu_buffer_type_get_alloc_size(ggml_backend_buffer
                     uint32_t       nwg = ggml_webgpu_flash_attn_vec_nwg(vec_nwg_cap, kv_tile, (uint32_t) K->ne[1]);
 
                     const size_t   align = capabilities.limits.minStorageBufferOffsetAlignment;
-                    const uint64_t nrows = (uint64_t) Q->ne[1] * Q->ne[2] * Q->ne[3];
+                    const uint64_t nrows = (uint64_t) n_q * Q->ne[2] * Q->ne[3];
                     if (nwg > 1u) {
                         const uint64_t tmp_data_elems  = nrows * (uint64_t) V->ne[0] * nwg;
                         const uint64_t tmp_stats_elems = nrows * 2u * nwg;
@@ -4036,7 +4068,7 @@ static size_t ggml_backend_webgpu_buffer_type_get_alloc_size(ggml_backend_buffer
                     }
                     if (mask != nullptr) {
                         const uint32_t blk_nblk0       = CEIL_DIV((uint32_t) K->ne[1], kv_tile);
-                        const uint32_t blk_nblk1       = CEIL_DIV((uint32_t) Q->ne[1], 1u);
+                        const uint32_t blk_nblk1       = (uint32_t) n_q;
                         const uint32_t stride_mask3    = (uint32_t) (mask->nb[3] / ggml_type_size(mask->type));
                         const uint32_t blk_batch_count = stride_mask3 > 0 ? (uint32_t) Q->ne[3] : 1u;
                         const uint64_t blk_elems       = (uint64_t) blk_nblk0 * blk_nblk1 * blk_batch_count;
@@ -4052,11 +4084,15 @@ static size_t ggml_backend_webgpu_buffer_type_get_alloc_size(ggml_backend_buffer
             {
                 const ggml_tensor * src0 = tensor->src[0];
                 const ggml_tensor * src1 = tensor->src[1];
-                bool                use_mmvq =
-                    ggml_webgpu_can_use_mmvq(src0, src1, ctx->webgpu_global_ctx->capabilities.supports_dot_product,
-                                             ctx->webgpu_global_ctx->vendor);
+
+                // the q8 copy of src1 grows with its rows, so a longer batch reserves the copy of the longest
+                // mmvq batch and the size stays monotonic in the batch length
+                const int64_t n_rows   = std::min<int64_t>(src1->ne[1], GGML_WEBGPU_MUL_MAT_VEC_MAX_ROWS);
+                const bool    use_mmvq = ggml_webgpu_can_use_mmvq(
+                    src0, src1, n_rows, ctx->webgpu_global_ctx->capabilities.supports_dot_product,
+                    ctx->webgpu_global_ctx->vendor);
                 if (use_mmvq) {
-                    const size_t q8_src1_size = src1->ne[3] * src1->ne[2] * src1->ne[1] *
+                    const size_t q8_src1_size = src1->ne[3] * src1->ne[2] * n_rows *
                                                 (36 /* sizeof(q8_1) */ * (src1->ne[0] / /* block_size */ 32));
                     res = ROUNDUP_POW2(res + q8_src1_size +
                                            ctx->webgpu_global_ctx->capabilities.limits.minStorageBufferOffsetAlignment,
@@ -4590,7 +4626,8 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
             {
                 // conservative support checks for whether the more resource-intensive shader paths
                 // can be used, to avoid cases where flash_attn is assigned to the CPU later on
-                supports_op = src0->type == GGML_TYPE_F32 &&
+                // K/V row lists (ggml_flash_attn_ext_rows) are not supported
+                supports_op = op->src[5] == nullptr && src0->type == GGML_TYPE_F32 &&
                               (src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16 ||
                                src1->type == GGML_TYPE_Q4_0 || src1->type == GGML_TYPE_Q8_0) &&
                               (src2->type == GGML_TYPE_F32 || src2->type == GGML_TYPE_F16 ||
@@ -4775,6 +4812,9 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
         case GGML_OP_COS:
             supports_op = (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) && (src0->type == op->type);
             break;
+        case GGML_OP_ARANGE:
+            supports_op = op->type == GGML_TYPE_F32 && ggml_is_contiguous(op);
+            break;
         case GGML_OP_PAD:
             supports_op = op->type == GGML_TYPE_F32 && src0->type == GGML_TYPE_F32;
             break;

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

@@ -0,0 +1,22 @@
+@group(0) @binding(0)
+var<storage, read_write> dst: array<f32>;
+
+struct Params {
+    ne: u32,
+    offset_dst: u32,
+    start: f32,
+    step: f32,
+};
+
+@group(0) @binding(1)
+var<uniform> params: Params;
+
+@compute @workgroup_size(WG_SIZE)
+fn main(@builtin(num_workgroups) num_wg: vec3<u32>, @builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x + (num_wg.x * u32(WG_SIZE)) * gid.y;
+    if (i >= params.ne) {
+        return;
+    }
+
+    dst[params.offset_dst + i] = params.start + params.step * f32(i);
+}

ggml/src/ggml-zdnn/common.hpp

@@ -34,6 +34,8 @@ struct ggml_backend_zdnn_device_context {
 
     size_t max_size;
 
+    int32_t n_threads;  // CPU threads
+
     char name[128];
 };
 

ggml/src/ggml-zdnn/ggml-zdnn.cpp

@@ -12,6 +12,7 @@
 #include <vector>
 #include <memory>
 #include <csignal>  // raise(SIGTRAP)
+#include <thread>
 #include <unistd.h>
 
 static void ggml_zdnn_compute_forward_mul_mat(
@@ -85,6 +86,15 @@ static enum ggml_status ggml_zdnn_graph_compute(ggml_backend_t backend, ggml_cgr
     GGML_UNUSED(ctx_dev);
 }
 
+static int64_t ggml_zdnn_mul_mat_min_batch(ggml_type type, int32_t n_threads) {
+    switch (type) {
+        case GGML_TYPE_Q8_0:
+            return std::max<int64_t>(1, n_threads - 2);
+        default:
+            return 1;
+    }
+}
+
 static bool ggml_zdnn_supports_op(const ggml_backend_zdnn_device_context * ctx_dev, const ggml_tensor * op) {
     switch (op->op) {
         case GGML_OP_NONE:
@@ -112,6 +122,11 @@ static bool ggml_zdnn_supports_op(const ggml_backend_zdnn_device_context * ctx_d
                         return false;
                 }
 
+                // offload smaller batches to the CPU
+                if (ne1 < ggml_zdnn_mul_mat_min_batch(weights->type, ctx_dev->n_threads)) {
+                    return false;
+                }
+
                 switch (weights->type) {
                     case GGML_TYPE_F32:
                     case GGML_TYPE_F16:
@@ -150,6 +165,7 @@ static ggml_backend_zdnn_device_context g_ggml_ctx_dev_main = {
     /* .has_parmblkformat_0   = */ false,
     /* .has_parmblkformat_1   = */ false,
     /* .max_size              = */ 0,
+    /* .n_threads             = */ 1,
     /* .name                  = */ "",
 };
 
@@ -164,6 +180,8 @@ static int ggml_backend_zdnn_device_acq(ggml_backend_zdnn_device_context * ctx)
         ctx->has_parmblkformat_0 = zdnn_is_nnpa_parmblk_fmt_installed(1, NNPA_PARMBLKFORMAT_0);
         ctx->has_parmblkformat_1 = zdnn_is_nnpa_parmblk_fmt_installed(1, NNPA_PARMBLKFORMAT_1);
         ctx->max_size = zdnn_get_nnpa_max_dim_idx_size();
+        // llama.cpp uses every core by default, and splits the first graph before it sets the threads
+        ctx->n_threads = std::max<int32_t>(1, (int32_t)std::thread::hardware_concurrency());
         strncpy(ctx->name, GGML_ZDNN_NAME, sizeof(ctx->name) - 1);
     }
 
@@ -452,8 +470,7 @@ static size_t ggml_backend_zdnn_buffer_type_get_alignment(ggml_backend_buffer_ty
 }
 
 static bool ggml_backend_zdnn_buffer_type_is_host(ggml_backend_buffer_type_t buft) {
-    /* while it resides in host memory, additional transformation is needed */
-    return false;
+    return true;
 
     GGML_UNUSED(buft);
 }
@@ -673,6 +690,9 @@ static void ggml_backend_zdnn_set_n_threads(ggml_backend_t backend, int n_thread
 
     ggml_backend_zdnn_context * ctx = (ggml_backend_zdnn_context *)backend->context;
     ctx->n_threads = n_threads;
+
+    ggml_backend_zdnn_device_context * ctx_dev = (ggml_backend_zdnn_device_context *)backend->device->context;
+    ctx_dev->n_threads = n_threads;
 }
 
 static void * ggml_backend_zdnn_get_proc_address(ggml_backend_reg_t reg, const char * name) {

ggml/src/ggml-zdnn/mmq.cpp

@@ -1,6 +1,7 @@
 #include "ggml.h"
 #include "mmq.hpp"
 #include "quantize.hpp"
+#include "utils.hpp"
 
 #include <algorithm>
 #include <future>
@@ -189,6 +190,11 @@ void ggml_zdnn_mul_mat_q(
     ggml_backend_zdnn_buffer * weights_extra = (ggml_backend_zdnn_buffer *)weights->extra;
     ggml_backend_zdnn_buffer * output_extra  = (ggml_backend_zdnn_buffer *)output->extra;
 
+    // weights read straight into a host buffer skip set_tensor, so transform them on first use
+    if (!weights_extra->ztensor.is_transformed) {
+        ggml_zdnn_load_tensor_q(weights_extra, weights);
+    }
+
     // starting threads costs more than it saves for small inputs, e.g. during token generation
     const int32_t n_threads = ne10*ne11 >= 65536 ? ctx->n_threads : 1;
 

ggml/src/ggml.c

@@ -71,12 +71,12 @@ uint64_t ggml_graph_next_uid(void) {
 
 // Needed for ggml_fp32_to_bf16_row()
 #if defined(__AVX512BF16__)
-#if defined(_MSC_VER)
+#if defined(_MSC_VER) && !defined(__clang__)
 #define m512i(p) p
 #else
 #include <immintrin.h>
 #define m512i(p) (__m512i)(p)
-#endif // defined(_MSC_VER)
+#endif // defined(_MSC_VER) && !defined(__clang__)
 #endif // defined(__AVX512BF16__)
 
 #if defined(__linux__) || \
@@ -357,7 +357,8 @@ void * ggml_aligned_malloc(size_t size) {
     void * aligned_memory = NULL;
   #ifdef GGML_USE_CPU_HBM
     int result = hbw_posix_memalign(&aligned_memory, alignment, size);
-  #elif TARGET_OS_OSX && !defined(GGML_TSAN_ENABLED)
+  #elif defined(__APPLE__) && !defined(GGML_TSAN_ENABLED)
+    // libmalloc can keep freed pages in the process footprint, vm_deallocate returns them at once
     GGML_UNUSED(alignment);
     kern_return_t alloc_status = vm_allocate((vm_map_t) mach_task_self(), (vm_address_t *) &aligned_memory, size, VM_FLAGS_ANYWHERE);
     int result = EFAULT;
@@ -404,7 +405,7 @@ void ggml_aligned_free(void * ptr, size_t size) {
     if (ptr != NULL) {
         hbw_free(ptr);
     }
-#elif TARGET_OS_OSX && !defined(GGML_TSAN_ENABLED)
+#elif defined(__APPLE__) && !defined(GGML_TSAN_ENABLED)
     if (ptr != NULL) {
         vm_deallocate((vm_map_t)mach_task_self(), (vm_address_t)ptr, size);
     }
@@ -1787,6 +1788,7 @@ static struct ggml_tensor * ggml_new_tensor_impl(
         enum   ggml_type      type,
         int                   n_dims,
         const int64_t       * ne,
+        const size_t        * nb,
         struct ggml_tensor  * view_src,
         size_t                view_offs) {
 
@@ -1813,8 +1815,6 @@ static struct ggml_tensor * ggml_new_tensor_impl(
         data_size *= ne[i];
     }
 
-    GGML_ASSERT(view_src == NULL || data_size == 0 || data_size + view_offs <= ggml_nbytes(view_src));
-
     void * data = view_src != NULL ? view_src->data : NULL;
     if (data != NULL) {
         data = (char *) data + view_offs;
@@ -1858,12 +1858,23 @@ static struct ggml_tensor * ggml_new_tensor_impl(
         result->ne[i] = ne[i];
     }
 
-    result->nb[0] = ggml_type_size(type);
-    result->nb[1] = result->nb[0]*(result->ne[0]/ggml_blck_size(type));
-    for (int i = 2; i < GGML_MAX_DIMS; i++) {
-        result->nb[i] = result->nb[i - 1]*result->ne[i - 1];
+    // a view built with strides takes them from the caller, any other tensor is contiguous
+    if (nb != NULL) {
+        for (int i = 0; i < GGML_MAX_DIMS; i++) {
+            result->nb[i] = nb[i];
+        }
+    } else {
+        result->nb[0] = ggml_type_size(type);
+        result->nb[1] = result->nb[0]*(result->ne[0]/ggml_blck_size(type));
+        for (int i = 2; i < GGML_MAX_DIMS; i++) {
+            result->nb[i] = result->nb[i - 1]*result->ne[i - 1];
+        }
     }
 
+    // a view stays inside its source, over the bytes its strides actually reach
+    const size_t view_size = ggml_nbytes(result);
+    GGML_ASSERT(view_src == NULL || view_size == 0 || view_size + view_offs <= ggml_nbytes(view_src));
+
     ctx->n_objects++;
 
     return result;
@@ -1874,7 +1885,7 @@ struct ggml_tensor * ggml_new_tensor(
         enum   ggml_type      type,
         int                   n_dims,
         const int64_t       * ne) {
-    return ggml_new_tensor_impl(ctx, type, n_dims, ne, NULL, 0);
+    return ggml_new_tensor_impl(ctx, type, n_dims, ne, NULL, NULL, 0);
 }
 
 struct ggml_tensor * ggml_new_tensor_1d(
@@ -1991,13 +2002,9 @@ struct ggml_tensor * ggml_format_name(struct ggml_tensor * tensor, const char *
 struct ggml_tensor * ggml_view_tensor(
         struct ggml_context * ctx,
         struct ggml_tensor  * src) {
-    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, src->type, GGML_MAX_DIMS, src->ne, src, 0);
+    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, src->type, GGML_MAX_DIMS, src->ne, src->nb, src, 0);
     ggml_format_name(result, "%s (view)", src->name);
 
-    for (int i = 0; i < GGML_MAX_DIMS; i++) {
-        result->nb[i] = src->nb[i];
-    }
-
     return result;
 }
 
@@ -2947,7 +2954,7 @@ static struct ggml_tensor * ggml_glu_impl(
     }
 
     int64_t ne[GGML_MAX_DIMS] = { a->ne[0] / 2 }; for (int i = 1; i < GGML_MAX_DIMS; i++) ne[i] = a->ne[i];
-    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, GGML_MAX_DIMS, b ? a->ne : ne, NULL, 0);
+    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, GGML_MAX_DIMS, b ? a->ne : ne, NULL, NULL, 0);
 
     ggml_set_op_params_i32(result, 0, (int32_t) op);
     ggml_set_op_params_i32(result, 1, (int32_t) swapped);
@@ -3756,7 +3763,7 @@ struct ggml_tensor * ggml_reshape(
     // as only the shape of b is relevant, and not its memory layout, b is allowed to be non contiguous.
     GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b));
 
-    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, GGML_MAX_DIMS, b->ne, a, 0);
+    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, GGML_MAX_DIMS, b->ne, NULL, a, 0);
     ggml_format_name(result, "%s (reshaped)", a->name);
 
     result->op     = GGML_OP_RESHAPE;
@@ -3773,7 +3780,7 @@ struct ggml_tensor * ggml_reshape_1d(
     GGML_ASSERT(ggml_nelements(a) == ne0);
 
     const int64_t ne[1] = { ne0 };
-    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 1, ne, a, 0);
+    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 1, ne, NULL, a, 0);
     ggml_format_name(result, "%s (reshaped)", a->name);
 
     result->op     = GGML_OP_RESHAPE;
@@ -3791,7 +3798,7 @@ struct ggml_tensor * ggml_reshape_2d(
     GGML_ASSERT(ggml_nelements(a) == ne0*ne1);
 
     const int64_t ne[2] = { ne0, ne1 };
-    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 2, ne, a, 0);
+    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 2, ne, NULL, a, 0);
     ggml_format_name(result, "%s (reshaped)", a->name);
 
     result->op     = GGML_OP_RESHAPE;
@@ -3810,7 +3817,7 @@ struct ggml_tensor * ggml_reshape_3d(
     GGML_ASSERT(ggml_nelements(a) == ne0*ne1*ne2);
 
     const int64_t ne[3] = { ne0, ne1, ne2 };
-    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 3, ne, a, 0);
+    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 3, ne, NULL, a, 0);
     ggml_format_name(result, "%s (reshaped)", a->name);
 
     result->op     = GGML_OP_RESHAPE;
@@ -3830,7 +3837,7 @@ struct ggml_tensor * ggml_reshape_4d(
     GGML_ASSERT(ggml_nelements(a) == ne0*ne1*ne2*ne3);
 
     const int64_t ne[4] = { ne0, ne1, ne2, ne3 };
-    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 4, ne, a, 0);
+    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 4, ne, NULL, a, 0);
     ggml_format_name(result, "%s (reshaped)", a->name);
 
     result->op     = GGML_OP_RESHAPE;
@@ -3844,8 +3851,9 @@ static struct ggml_tensor * ggml_view_impl(
         struct ggml_tensor  * a,
         int                   n_dims,
         const int64_t       * ne,
+        const size_t        * nb,
         size_t                offset) {
-    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, n_dims, ne, a, offset);
+    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, n_dims, ne, nb, a, offset);
     ggml_format_name(result, "%s (view)", a->name);
 
     ggml_set_op_params(result, &offset, sizeof(offset));
@@ -3863,7 +3871,7 @@ struct ggml_tensor * ggml_view_1d(
         struct ggml_tensor  * a,
         int64_t               ne0,
         size_t                offset) {
-    struct ggml_tensor * result = ggml_view_impl(ctx, a, 1, &ne0, offset);
+    struct ggml_tensor * result = ggml_view_impl(ctx, a, 1, &ne0, NULL, offset);
 
     return result;
 }
@@ -3878,12 +3886,9 @@ struct ggml_tensor * ggml_view_2d(
         size_t                nb1,
         size_t                offset) {
     const int64_t ne[2] = { ne0, ne1 };
+    const size_t  nb[4] = { ggml_type_size(a->type), nb1, nb1*ne1, nb1*ne1 };
 
-    struct ggml_tensor * result = ggml_view_impl(ctx, a, 2, ne, offset);
-
-    result->nb[1] = nb1;
-    result->nb[2] = result->nb[1]*ne1;
-    result->nb[3] = result->nb[2];
+    struct ggml_tensor * result = ggml_view_impl(ctx, a, 2, ne, nb, offset);
 
     return result;
 }
@@ -3900,12 +3905,9 @@ struct ggml_tensor * ggml_view_3d(
         size_t                nb2,
         size_t                offset) {
     const int64_t ne[3] = { ne0, ne1, ne2 };
+    const size_t  nb[4] = { ggml_type_size(a->type), nb1, nb2, nb2*ne2 };
 
-    struct ggml_tensor * result = ggml_view_impl(ctx, a, 3, ne, offset);
-
-    result->nb[1] = nb1;
-    result->nb[2] = nb2;
-    result->nb[3] = result->nb[2]*ne2;
+    struct ggml_tensor * result = ggml_view_impl(ctx, a, 3, ne, nb, offset);
 
     return result;
 }
@@ -3924,12 +3926,9 @@ struct ggml_tensor * ggml_view_4d(
         size_t                nb3,
         size_t                offset) {
     const int64_t ne[4] = { ne0, ne1, ne2, ne3 };
+    const size_t  nb[4] = { ggml_type_size(a->type), nb1, nb2, nb3 };
 
-    struct ggml_tensor * result = ggml_view_impl(ctx, a, 4, ne, offset);
-
-    result->nb[1] = nb1;
-    result->nb[2] = nb2;
-    result->nb[3] = nb3;
+    struct ggml_tensor * result = ggml_view_impl(ctx, a, 4, ne, nb, offset);
 
     return result;
 }
@@ -5388,10 +5387,10 @@ struct ggml_tensor * ggml_roll(
         int                   shift2,
         int                   shift3) {
     GGML_ASSERT(a->nb[0] == ggml_type_size(a->type));
-    GGML_ASSERT(abs(shift0) < a->ne[0]);
-    GGML_ASSERT(abs(shift1) < a->ne[1]);
-    GGML_ASSERT(abs(shift2) < a->ne[2]);
-    GGML_ASSERT(abs(shift3) < a->ne[3]);
+    GGML_ASSERT(shift0 > -a->ne[0] && shift0 < a->ne[0]);
+    GGML_ASSERT(shift1 > -a->ne[1] && shift1 < a->ne[1]);
+    GGML_ASSERT(shift2 > -a->ne[2] && shift2 < a->ne[2]);
+    GGML_ASSERT(shift3 > -a->ne[3] && shift3 < a->ne[3]);
 
     struct ggml_tensor * result = ggml_dup_tensor(ctx, a);
 
@@ -5556,20 +5555,30 @@ struct ggml_tensor * ggml_arange(
 
 // ggml_flash_attn_ext
 
-struct ggml_tensor * ggml_flash_attn_ext(
+static struct ggml_tensor * ggml_flash_attn_ext_impl(
         struct ggml_context * ctx,
         struct ggml_tensor  * q,
         struct ggml_tensor  * k,
         struct ggml_tensor  * v,
         struct ggml_tensor  * mask,
+        struct ggml_tensor  * kv_rows,
         float                 scale,
         float                 max_bias,
         float                 logit_softcap) {
     GGML_ASSERT(ggml_can_mul_mat(k, q));
     // TODO: check if vT can be multiplied by (k*qT)
 
-    GGML_ASSERT(q->ne[3] == k->ne[3]);
-    GGML_ASSERT(q->ne[3] == v->ne[3]);
+    if (kv_rows) {
+        // k and v are shared by all q slices, kv_rows picks the rows of each slice
+        GGML_ASSERT(mask && mask->ne[0] == kv_rows->ne[0]);
+        GGML_ASSERT(kv_rows->type == GGML_TYPE_I32 && ggml_is_contiguous(kv_rows));
+        GGML_ASSERT(kv_rows->ne[1] == q->ne[3] && kv_rows->ne[2] == 1 && kv_rows->ne[3] == 1);
+        GGML_ASSERT(k->ne[3] == 1);
+        GGML_ASSERT(v->ne[3] == 1);
+    } else {
+        GGML_ASSERT(q->ne[3] == k->ne[3]);
+        GGML_ASSERT(q->ne[3] == v->ne[3]);
+    }
 
     if (mask) {
         GGML_ASSERT(mask->type == GGML_TYPE_F16);
@@ -5596,10 +5605,36 @@ struct ggml_tensor * ggml_flash_attn_ext(
     result->src[1] = k;
     result->src[2] = v;
     result->src[3] = mask;
+    result->src[5] = kv_rows;
 
     return result;
 }
 
+struct ggml_tensor * ggml_flash_attn_ext(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * q,
+        struct ggml_tensor  * k,
+        struct ggml_tensor  * v,
+        struct ggml_tensor  * mask,
+        float                 scale,
+        float                 max_bias,
+        float                 logit_softcap) {
+    return ggml_flash_attn_ext_impl(ctx, q, k, v, mask, NULL, scale, max_bias, logit_softcap);
+}
+
+struct ggml_tensor * ggml_flash_attn_ext_rows(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * q,
+        struct ggml_tensor  * k,
+        struct ggml_tensor  * v,
+        struct ggml_tensor  * mask,
+        struct ggml_tensor  * kv_rows,
+        float                 scale,
+        float                 max_bias,
+        float                 logit_softcap) {
+    return ggml_flash_attn_ext_impl(ctx, q, k, v, mask, kv_rows, scale, max_bias, logit_softcap);
+}
+
 
 void ggml_flash_attn_ext_set_prec(
         struct ggml_tensor * a,

gguf-py/gguf/constants.py

@@ -506,6 +506,11 @@ class Keys:
         WINDOW_SIZE         = "clip.audio.window_size"
         LOCAL_BLOCK_COUNT   = "clip.audio.local_block_count" # mimo-v2.5: input_local_transformer layer count
         LOCAL_GROUP_SIZE    = "clip.audio.local_group_size"  # mimo-v2.5: input_local_transformer grouping size
+        ROPE_DIMS           = "clip.audio.rope_dims"         # nemotron-transcribe
+
+        class Speaker: # nemotron-transcribe
+            BLOCK_COUNT         = "clip.audio.spk.block_count"
+            HEAD_COUNT          = "clip.audio.spk.head_count"
 
         class Attention:
             HEAD_COUNT      = "clip.audio.attention.head_count"
@@ -615,6 +620,7 @@ class MODEL_ARCH(IntEnum):
     QWEN35MOE        = auto()
     CLEF             = auto()
     VLLM_DECISION2   = auto()
+    CLOVER_MOE       = auto()
     QWEN4EXP         = auto()
     PHI2             = auto()
     PHI3             = auto()
@@ -1410,6 +1416,24 @@ class MODEL_TENSOR(IntEnum):
     A_ENC_CONV_NORM_VAR    = auto() # parakeet
     A_ENC_MEL_FILTERS      = auto() # parakeet
     A_ENC_WINDOW           = auto() # parakeet
+    A_ENC_ATTN_QKV         = auto() # nemotron-transcribe
+    A_SPK_INP_PROJ         = auto() # nemotron-transcribe
+    A_SPK_PRE_NORM         = auto() # nemotron-transcribe
+    A_SPK_POST_NORM        = auto() # nemotron-transcribe
+    A_SPK_ATTN_NORM        = auto() # nemotron-transcribe
+    A_SPK_ATTN_QKV         = auto() # nemotron-transcribe
+    A_SPK_ATTN_OUT         = auto() # nemotron-transcribe
+    A_SPK_FFN_NORM         = auto() # nemotron-transcribe
+    A_SPK_FFN_UP           = auto() # nemotron-transcribe
+    A_SPK_FFN_DOWN         = auto() # nemotron-transcribe
+    A_SPK_HEAD_PROJ        = auto() # nemotron-transcribe
+    A_SPK_HEAD_UPSAMPLE    = auto() # nemotron-transcribe
+    A_SPK_HEAD_FC1         = auto() # nemotron-transcribe
+    A_SPK_HEAD_FC2         = auto() # nemotron-transcribe
+    A_SPK_SIL_EMB          = auto() # nemotron-transcribe
+    A_FUSION_ASR_NORM      = auto() # nemotron-transcribe
+    A_FUSION_SPK_NORM      = auto() # nemotron-transcribe
+    A_FUSION_SPK_KERNEL    = auto() # nemotron-transcribe
     A_CTC_OUT              = auto()
     A_CTC_OUT_MID          = auto()
     A_ENC_ATTN_REL_POS_EMB = auto()
@@ -1473,6 +1497,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
     MODEL_ARCH.QWEN35MOE:        "qwen35moe",
     MODEL_ARCH.CLEF:             "clef",
     MODEL_ARCH.VLLM_DECISION2:   "vllm-decision2",
+    MODEL_ARCH.CLOVER_MOE:       "clover_moe",
     MODEL_ARCH.QWEN4EXP:         "qwen4exp",
     MODEL_ARCH.PHI2:             "phi2",
     MODEL_ARCH.PHI3:             "phi3",
@@ -2217,6 +2242,24 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
     MODEL_TENSOR.A_ENC_CONV_NORM_VAR:       "a.blk.{bid}.conv_norm_var",
     MODEL_TENSOR.A_ENC_MEL_FILTERS:         "a.mel_filters",
     MODEL_TENSOR.A_ENC_WINDOW:              "a.window",
+    MODEL_TENSOR.A_ENC_ATTN_QKV:            "a.blk.{bid}.attn_qkv",
+    MODEL_TENSOR.A_SPK_INP_PROJ:            "a.spk.input_projection",
+    MODEL_TENSOR.A_SPK_PRE_NORM:            "a.spk.pre_ln",
+    MODEL_TENSOR.A_SPK_POST_NORM:           "a.spk.post_ln",
+    MODEL_TENSOR.A_SPK_ATTN_NORM:           "a.spk.blk.{bid}.ln1",
+    MODEL_TENSOR.A_SPK_ATTN_QKV:            "a.spk.blk.{bid}.attn_qkv",
+    MODEL_TENSOR.A_SPK_ATTN_OUT:            "a.spk.blk.{bid}.attn_out",
+    MODEL_TENSOR.A_SPK_FFN_NORM:            "a.spk.blk.{bid}.ln2",
+    MODEL_TENSOR.A_SPK_FFN_UP:              "a.spk.blk.{bid}.ffn_up",
+    MODEL_TENSOR.A_SPK_FFN_DOWN:            "a.spk.blk.{bid}.ffn_down",
+    MODEL_TENSOR.A_SPK_HEAD_PROJ:           "a.spk.head.proj",
+    MODEL_TENSOR.A_SPK_HEAD_UPSAMPLE:       "a.spk.head.upsample",
+    MODEL_TENSOR.A_SPK_HEAD_FC1:            "a.spk.head.fc1",
+    MODEL_TENSOR.A_SPK_HEAD_FC2:            "a.spk.head.fc2",
+    MODEL_TENSOR.A_SPK_SIL_EMB:             "a.spk.head.sil_emb",
+    MODEL_TENSOR.A_FUSION_ASR_NORM:         "a.fusion.asr_norm",
+    MODEL_TENSOR.A_FUSION_SPK_NORM:         "a.fusion.spk_norm",
+    MODEL_TENSOR.A_FUSION_SPK_KERNEL:       "a.fusion.spk_kernel",
     MODEL_TENSOR.A_CTC_OUT:                 "a.enc_ctc_out",
     MODEL_TENSOR.A_CTC_OUT_MID:             "a.enc_ctc_out_mid",
     MODEL_TENSOR.A_ENC_ATTN_REL_POS_EMB:    "a.blk.{bid}.attn_rel_pos_emb",
@@ -2585,6 +2628,24 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.A_ENC_CONV_NORM_VAR,
         MODEL_TENSOR.A_ENC_MEL_FILTERS,
         MODEL_TENSOR.A_ENC_WINDOW,
+        MODEL_TENSOR.A_ENC_ATTN_QKV,
+        MODEL_TENSOR.A_SPK_INP_PROJ,
+        MODEL_TENSOR.A_SPK_PRE_NORM,
+        MODEL_TENSOR.A_SPK_POST_NORM,
+        MODEL_TENSOR.A_SPK_ATTN_NORM,
+        MODEL_TENSOR.A_SPK_ATTN_QKV,
+        MODEL_TENSOR.A_SPK_ATTN_OUT,
+        MODEL_TENSOR.A_SPK_FFN_NORM,
+        MODEL_TENSOR.A_SPK_FFN_UP,
+        MODEL_TENSOR.A_SPK_FFN_DOWN,
+        MODEL_TENSOR.A_SPK_HEAD_PROJ,
+        MODEL_TENSOR.A_SPK_HEAD_UPSAMPLE,
+        MODEL_TENSOR.A_SPK_HEAD_FC1,
+        MODEL_TENSOR.A_SPK_HEAD_FC2,
+        MODEL_TENSOR.A_SPK_SIL_EMB,
+        MODEL_TENSOR.A_FUSION_ASR_NORM,
+        MODEL_TENSOR.A_FUSION_SPK_NORM,
+        MODEL_TENSOR.A_FUSION_SPK_KERNEL,
         MODEL_TENSOR.A_MM_INP_PROJ,
         MODEL_TENSOR.A_MM_SOFT_EMB_NORM,
         MODEL_TENSOR.A_MM_EMBEDDING,
@@ -3383,6 +3444,37 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
         MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
     ],
+    MODEL_ARCH.CLOVER_MOE: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_INP_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_GATE_UP_EXP,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_CONV1D,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_NORM,
+        MODEL_TENSOR.SSM_BETA,
+        MODEL_TENSOR.SSM_ALPHA,
+        MODEL_TENSOR.SSM_OUT,
+    ],
     MODEL_ARCH.QWEN4EXP: [
         MODEL_TENSOR.TOKEN_EMBD,
         MODEL_TENSOR.OUTPUT,
@@ -6707,6 +6799,7 @@ class VisionProjectorType:
     MINICPMV4_7    = "minicpmv4_7"
     MINICPM_OCR_1_0 = "minicpm_ocr_1_0"
     GRANITE_SPEECH = "granite_speech"  # audio
+    NEMOTRON_TRANSCRIBE = "nemotron_transcribe"  # audio
     MIMOVL         = "mimovl"
     MIMO_AUDIO     = "mimo_audio"
     GRANITE4_VISION = "granite4_vision"

gguf-py/gguf/gguf_writer.py

@@ -1712,6 +1712,15 @@ class GGUFWriter:
     def add_audio_projector_head_count(self, value: int) -> None:
         self.add_uint32(Keys.ClipAudio.Projector.HEAD_COUNT, value)
 
+    def add_audio_rope_dims(self, value: int) -> None:
+        self.add_uint32(Keys.ClipAudio.ROPE_DIMS, value)
+
+    def add_audio_spk_block_count(self, value: int) -> None:
+        self.add_uint32(Keys.ClipAudio.Speaker.BLOCK_COUNT, value)
+
+    def add_audio_spk_head_count(self, value: int) -> None:
+        self.add_uint32(Keys.ClipAudio.Speaker.HEAD_COUNT, value)
+
     # audio generation (mmproj)
 
     def add_clip_gen_audio_projector_type(self, value: str) -> None:

gguf-py/gguf/tensor_mapping.py

@@ -2532,6 +2532,7 @@ class TensorNameMap:
 
         MODEL_TENSOR.A_ENC_INP_PROJ: (
             "conformer.subsample_conv_projection.input_proj_linear", # gemma4
+            "perception.encoder.asr_encoder.pre_encode.proj", # nemotron-transcribe (feature stacking)
             "encoder.input_linear",
         ),
 
@@ -2546,13 +2547,16 @@ class TensorNameMap:
             "speaker_encoder.mfa.conv", # qwen3tts speaker encoder: multi-layer feature aggregation
         ),
 
-        MODEL_TENSOR.A_PRE_NORM: (),
+        MODEL_TENSOR.A_PRE_NORM: (
+            "perception.encoder.asr_encoder.embed_norm", # nemotron-transcribe
+        ),
 
         MODEL_TENSOR.A_POST_NORM: (
             "audio_tower.layer_norm", # ultravox
             "audio_tower.norm", # voxtral realtime
             "audio_encoder.dots_encoder.speech_encoder.layer_norm", # dots3note
             "audio_tower.ln_post", # qwen2omni
+            "perception.encoder.asr_encoder.final_norm", # nemotron-transcribe
             "encoder.layer_norm", # mimo-audio-tokenizer
             "encoder.lfm.lfm.embedding_norm", # lfm3a
         ),
@@ -2595,10 +2599,12 @@ class TensorNameMap:
 
         MODEL_TENSOR.A_ENC_ATTN_Q_NORM: (
             "encoder.lfm.lfm.layers.{bid}.self_attn.q_layernorm", # lfm3a
+            "perception.encoder.asr_encoder.layers.{bid}.attn.q_norm", # nemotron-transcribe
         ),
 
         MODEL_TENSOR.A_ENC_ATTN_K_NORM: (
             "encoder.lfm.lfm.layers.{bid}.self_attn.k_layernorm", # lfm3a
+            "perception.encoder.asr_encoder.layers.{bid}.attn.k_norm", # nemotron-transcribe
         ),
 
         MODEL_TENSOR.A_ENC_ATTN_K_REL: (
@@ -2629,6 +2635,7 @@ class TensorNameMap:
             "encoder.lfm.lfm.layers.{bid}.operator_norm", # lfm3a
             "conformer.layers.{bid}.attention.pre_attn_norm", # gemma3n
             "sound_encoder.encoder.layers.{bid}.norm_self_att", # parakeet
+            "perception.encoder.asr_encoder.layers.{bid}.norm1", # nemotron-transcribe
             "encoder.layers.{bid}.attn.pre_norm", # granite_speech
             "encoder.layers.{bid}.self_attn_layer_norm", # mimo-audio-tokenizer
         ),
@@ -2642,6 +2649,7 @@ class TensorNameMap:
             "conformer.layers.{bid}.attention.post", # gemma3n
             "conformer.layers.{bid}.self_attn.post", # gemma4
             "sound_encoder.encoder.layers.{bid}.self_attn.o_proj", # parakeet
+            "perception.encoder.asr_encoder.layers.{bid}.attn.out_proj", # nemotron-transcribe
             "encoder.layers.{bid}.attn.to_out", # granite_speech
             "encoder.layers.{bid}.self_attn.out_proj", # mimo-audio-tokenizer
         ),
@@ -2653,6 +2661,7 @@ class TensorNameMap:
             "encoder.lfm.lfm.layers.{bid}.ffn_norm", # lfm3a
             "conformer.layers.{bid}.attention.post_norm", # gemma3n
             "sound_encoder.encoder.layers.{bid}.norm_out", # parakeet
+            "perception.encoder.asr_encoder.layers.{bid}.norm2", # nemotron-transcribe
             "encoder.layers.{bid}.post_norm", # granite_speech
             "encoder.layers.{bid}.final_layer_norm", # mimo-audio-tokenizer
         ),
@@ -2683,6 +2692,7 @@ class TensorNameMap:
             "conformer.layers.{bid}.ffw_layer_start.ffw_layer_1", # gemma3n
             "conformer.layers.{bid}.feed_forward1.ffw_layer_1", # gemma4
             "sound_encoder.encoder.layers.{bid}.feed_forward1.linear1", # parakeet
+            "perception.encoder.asr_encoder.layers.{bid}.ffn.net.0", # nemotron-transcribe
             "encoder.layers.{bid}.ff1.up_proj", # granite_speech
             "encoder.layers.{bid}.fc1", # mimo-audio-tokenizer
         ),
@@ -2702,6 +2712,7 @@ class TensorNameMap:
             "conformer.layers.{bid}.ffw_layer_start.ffw_layer_2", # gemma3n
             "conformer.layers.{bid}.feed_forward1.ffw_layer_2", # gemma4
             "sound_encoder.encoder.layers.{bid}.feed_forward1.linear2", # parakeet
+            "perception.encoder.asr_encoder.layers.{bid}.ffn.net.3", # nemotron-transcribe
             "encoder.layers.{bid}.ff1.down_proj", # granite_speech
             "encoder.layers.{bid}.fc2", # mimo-audio-tokenizer
         ),
@@ -2792,6 +2803,7 @@ class TensorNameMap:
             "audio_tower.proj", # qwen2omni
             "model.audio_tower.output_proj", # gemma4
             "speaker_encoder.fc", # qwen3tts speaker encoder: final speaker embedding projection
+            "perception.proj", # nemotron-transcribe
         ),
 
         MODEL_TENSOR.A_MM_NORM_PRE: (
@@ -2914,12 +2926,85 @@ class TensorNameMap:
 
         MODEL_TENSOR.A_ENC_MEL_FILTERS: (
             "sound_encoder.encoder.feature_extractor.featurizer.fb", # parakeet
+            "perception.preprocessor.featurizer.fb", # nemotron-transcribe
         ),
 
         MODEL_TENSOR.A_ENC_WINDOW: (
             "sound_encoder.encoder.feature_extractor.featurizer.window", # parakeet
+            "perception.preprocessor.featurizer.window", # nemotron-transcribe
         ),
 
+        MODEL_TENSOR.A_ENC_ATTN_QKV: (
+            "perception.encoder.asr_encoder.layers.{bid}.attn.w_qkv", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_INP_PROJ: (
+            "perception.encoder.diarization_model.encoder.pre_encode.proj", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_PRE_NORM: (
+            "perception.encoder.diarization_model.encoder.embed_norm", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_POST_NORM: (
+            "perception.encoder.diarization_model.encoder.final_norm", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_ATTN_NORM: (
+            "perception.encoder.diarization_model.encoder.layers.{bid}.norm1", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_ATTN_QKV: (
+            "perception.encoder.diarization_model.encoder.layers.{bid}.attn.w_qkv", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_ATTN_OUT: (
+            "perception.encoder.diarization_model.encoder.layers.{bid}.attn.out_proj", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_FFN_NORM: (
+            "perception.encoder.diarization_model.encoder.layers.{bid}.norm2", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_FFN_UP: (
+            "perception.encoder.diarization_model.encoder.layers.{bid}.ffn.net.0", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_FFN_DOWN: (
+            "perception.encoder.diarization_model.encoder.layers.{bid}.ffn.net.3", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_HEAD_PROJ: (
+            "perception.encoder.diarization_model.sortformer_modules.encoder_proj", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_HEAD_UPSAMPLE: (
+            "perception.encoder.diarization_model.sortformer_modules.subpixel_upsample", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_HEAD_FC1: (
+            "perception.encoder.diarization_model.sortformer_modules.first_hidden_to_hidden", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_HEAD_FC2: (
+            "perception.encoder.diarization_model.sortformer_modules.single_hidden_to_spks", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_SIL_EMB: (
+            "perception.encoder.diarization_model.sortformer_modules.learnable_sil_emb", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_FUSION_ASR_NORM: (
+            "perception.encoder.asr_norm", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_FUSION_SPK_NORM: (
+            "perception.encoder.diar_norm", # nemotron-transcribe
+        ),
+
+        # generated in conversion code (nemotron-transcribe)
+        MODEL_TENSOR.A_FUSION_SPK_KERNEL: (),
+
         MODEL_TENSOR.A_MM_EMBEDDING: (
             "model.embed_audio.embedding", # gemma3n
         ),