qwen4exp: fix sparse-attention block selection

Build QSA blocks per sequence in token order and select complete blocks
before expanding them to cache cells. Keep only the incomplete tail
unconditionally visible and rotate pooled keys with the first token's
full M-RoPE position.

This prevents unified-cache sequences from sharing pooled indexer keys
and avoids replacing padded tail entries with extra history tokens.

Synthetic Qwen4 architecture, exact mask, F16 and Q8_0 state,
sequence-copy, Metal, and AddressSanitizer checks pass.

Assisted-by: Codex

qwen4exp: support independent PLE embedding widths

Size the PLE key and value projections from the concatenated n-gram
embedding instead of assuming it matches the model hidden width.
Validate the head count before narrowing it to the stored type.

Add a synthetic PLE model with a 64-wide embedding and a 256-wide
hidden state, then verify inference and model roundtrip.

Assisted-by: Codex

qwen4exp: validate model metadata

Reject invalid GDN, hyper-connection, QSA, and PLE dimensions during
model loading instead of aborting later while building the graph.
Validate PLE array lengths before copying them into fixed storage.

The released configuration and synthetic Qwen4 architecture tests pass.

Assisted-by: Codex

qwen4exp: update indexer cache after sequence copies

Treat cached indexer keys as unrotated data and apply pending cache
updates alongside the attention and recurrent state. This copies indexer
data during non-unified cross-stream sequence copies without applying
RoPE shifts to raw keys.

Assisted-by: Codex

qwen4exp: enable recurrent state rollback

Assisted-by: Codex

qwen4exp: disable tensor split

Assisted-by: Codex

metal: align dynamic threadgroup memory

Assisted-by: Codex

metal: widen expert matmul thread index

Assisted-by: Codex

common/arg.cpp

@@ -2735,9 +2735,9 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         "- auto: on, but only for tensors larger than 4 GiB\n"
         "- off: always keep them resident",
         [](common_params & params, const std::string & value) {
-            /**/ if (value == "on")   { params.lazy_mode = LLAMA_LAZY_MODE_ON;   }
-            else if (value == "auto") { params.lazy_mode = LLAMA_LAZY_MODE_AUTO; }
-            else if (value == "off")  { params.lazy_mode = LLAMA_LAZY_MODE_OFF;  }
+            /**/ if (value == "on")   { params.tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_ON;   }
+            else if (value == "auto") { params.tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_AUTO; }
+            else if (value == "off")  { params.tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_OFF;  }
             else { throw std::invalid_argument("invalid value"); }
         }
     ).set_env("LLAMA_ARG_TENSOR_READ_LAZY"));

common/common.cpp

@@ -1688,7 +1688,7 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
     mparams.main_gpu        = params.main_gpu;
     mparams.split_mode      = params.split_mode;
     mparams.load_mode       = params.load_mode;
-    mparams.lazy_mode = params.lazy_mode;
+    mparams.tensor_read_lazy = params.tensor_read_lazy;
     mparams.tensor_split    = params.tensor_split;
     mparams.check_tensors   = params.check_tensors;
     mparams.use_extra_bufts = !params.no_extra_bufts;

common/common.h

@@ -483,7 +483,7 @@ struct common_params {
     enum llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER; // how to split the model across GPUs
     enum llama_load_mode  load_mode  = LLAMA_LOAD_MODE_AUTO; // how to load the model
 
-    enum llama_lazy_mode lazy_mode = LLAMA_LAZY_MODE_AUTO; // on-demand reading of tensors marked by the arch
+    enum llama_tensor_read_lazy tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_AUTO; // on-demand reading of tensors marked by the arch
 
     common_cpu_params cpuparams;
     common_cpu_params cpuparams_batch;

docs/backend/snapdragon/windows.md

@@ -24,18 +24,7 @@ must be included in the .cat file digitally signed with a trusted certificate.
 This document covers details on how to generate personal certificate files (.pfx) and how to configure the system
 to allow for test signatures (aka test-signing).
 
-## Install Windows SDKs
-
-The recommended method is `setup-sdk.py`:
-
-```
-> python scripts\snapdragon\setup-sdk.py --list-sdk-releases
-> python scripts\snapdragon\setup-sdk.py --hexagon --opencl
-```
-
-It installs the selected SDKs under `C:\Qualcomm` and sets their corresponding environment variables for the current user. Start a new terminal after it completes; native Windows builds check all SDK paths before CMake runs.
-
-Select the SDKs to install with `--hexagon` and `--opencl`; use both to prepare a dual-backend build. To select a different available version, pass it to the SDK option, for example `--hexagon 6.4.0.2`. SDK versions install side by side, so you can switch versions without deleting an existing installation. Use `--force` to reinstall the selected SDKs. Use a new CMake build directory after each switch because CMake caches the SDK paths.
+## Install the latest Adreno OpenCL SDK
 
 Either use the trimmed down version (optimized for CI) from
 

ggml/src/ggml-backend-impl.h

@@ -103,16 +103,6 @@ extern "C" {
     // Backend (stream)
     //
 
-    // passed to graph_optimize so the backend can add allocation dependencies:
-    // if the backend executes parts of the graph out of order (e.g. on concurrent streams),
-    // it must keep the affected tensors allocated until a node where execution is known to have joined
-    struct ggml_backend_graph_optimize_params {
-        // keep `tensor` allocated at least until `until` (a node of the same graph) has been computed
-        // can be called multiple times for the same tensor: the longest lifetime applies
-        void (*add_alloc_dep)(void * user_data, struct ggml_tensor * tensor, struct ggml_tensor * until);
-        void * user_data;
-    };
-
     struct ggml_backend_i {
         const char * (*get_name)(ggml_backend_t backend);
 
@@ -147,7 +137,7 @@ extern "C" {
         void (*event_wait)  (ggml_backend_t backend, ggml_backend_event_t event);
 
         // (optional) sort/optimize the nodes in the graph
-        void                      (*graph_optimize)    (ggml_backend_t backend, struct ggml_cgraph * cgraph, struct ggml_backend_graph_optimize_params * params);
+        void                      (*graph_optimize)    (ggml_backend_t backend, struct ggml_cgraph * cgraph);
     };
 
     struct ggml_backend {

ggml/src/ggml-backend.cpp

@@ -20,7 +20,6 @@
 #include <stdlib.h>
 #include <string.h>
 #include <algorithm>
-#include <unordered_map>
 #include <vector>
 
 #ifdef __APPLE__
@@ -559,10 +558,10 @@ void ggml_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t event)
     backend->iface.event_wait(backend, event);
 }
 
-static void ggml_backend_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * cgraph, struct ggml_backend_graph_optimize_params * params) {
+static void ggml_backend_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * cgraph) {
     GGML_ASSERT(backend);
     if (backend->iface.graph_optimize != NULL) {
-        backend->iface.graph_optimize(backend, cgraph, params);
+        backend->iface.graph_optimize(backend, cgraph);
     }
 }
 
@@ -1442,40 +1441,11 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra
         sched->prev_leaf_backend_ids = tmp;
     }
 
-    // optimize the split graphs and collect the allocation dependencies added by the backends
-    // this needs to happen before we make graph_copy, so they are in sync
-    // TODO: this may create many small allocations in the scheduler, restructure to use a flat array
-    std::unordered_map<ggml_tensor *, std::vector<ggml_tensor *>> alloc_deps;
-
-    struct ggml_backend_graph_optimize_params opt_params = {
-        /* .add_alloc_dep = */ [](void * user_data, ggml_tensor * tensor, ggml_tensor * until) {
-            auto & deps = *(std::unordered_map<ggml_tensor *, std::vector<ggml_tensor *>> *) user_data;
-            std::vector<ggml_tensor *> & keep = deps[until];
-            if (std::find(keep.begin(), keep.end(), tensor) == keep.end()) {
-                keep.push_back(tensor);
-            }
-        },
-        /* .user_data     = */ &alloc_deps,
-    };
-
-    for (int i = 0; i < sched->n_splits; i++) {
-        struct ggml_backend_sched_split * split = &sched->splits[i];
-        split->graph = ggml_graph_view(graph, split->i_start, split->i_end);
-
-        ggml_backend_graph_optimize(sched->backends[split->backend_id], &split->graph, &opt_params);
-    }
-
-    // each dep is added to graph_copy as a GGML_OP_NONE node with the kept tensors as srcs
-    int n_dep_nodes = 0;
-    for (const auto & it : alloc_deps) {
-        n_dep_nodes += (it.second.size() + GGML_MAX_SRC - 1) / GGML_MAX_SRC;
-    }
-
     int total_inputs = sched->n_graph_inputs;
     for (int i = 0; i < sched->n_splits; i++) {
         total_inputs += sched->splits[i].n_inputs;
     }
-    int graph_size = std::max(graph->n_nodes, graph->n_leafs) + total_inputs * 2 * sched->n_copies + n_dep_nodes;
+    int graph_size = std::max(graph->n_nodes, graph->n_leafs) + total_inputs * 2 * sched->n_copies;
 
     // remember the actual graph_size for performing reallocation checks later [GGML_SCHED_DEBUG_REALLOC]
     sched->debug_prev_graph_size = sched->debug_graph_size;
@@ -1493,10 +1463,13 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra
 
     struct ggml_cgraph * graph_copy = &sched->graph;
 
-    int n_dep_nodes_added = 0;
-
     for (int i = 0; i < sched->n_splits; i++) {
         struct ggml_backend_sched_split * split = &sched->splits[i];
+        split->graph = ggml_graph_view(graph, split->i_start, split->i_end);
+
+        // Optimize this split of the graph. This needs to happen before we make graph_copy,
+        // so they are in sync.
+        ggml_backend_graph_optimize(sched->backends[split->backend_id], &split->graph);
 
         // add inputs to the graph copy so that they are allocated by ggml-alloc at the start of the split
         for (int j = 0; j < split->n_inputs; j++) {
@@ -1521,32 +1494,9 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra
             assert(graph_copy->size > graph_copy->n_nodes);
             sched->node_backend_ids[graph_copy->n_nodes] = tensor_backend_id(graph->nodes[j]);
             graph_copy->nodes[graph_copy->n_nodes++] = graph->nodes[j];
-
-            if (alloc_deps.empty()) {
-                continue;
-            }
-
-            // add a dependency node so that the kept tensors are not freed before this node is computed
-            auto it = alloc_deps.find(graph->nodes[j]);
-            if (it != alloc_deps.end()) {
-                const std::vector<ggml_tensor *> & keep = it->second;
-                for (size_t k = 0; k < keep.size(); k += GGML_MAX_SRC) {
-                    struct ggml_tensor * dep = ggml_view_tensor(sched->ctx, keep[k]);
-                    for (size_t s = 0; s < GGML_MAX_SRC && k + s < keep.size(); s++) {
-                        dep->src[s] = keep[k + s];
-                    }
-                    assert(graph_copy->size > graph_copy->n_nodes);
-                    sched->node_backend_ids[graph_copy->n_nodes] = split->backend_id;
-                    graph_copy->nodes[graph_copy->n_nodes++] = dep;
-                    n_dep_nodes_added++;
-                }
-            }
         }
     }
 
-    // a mismatch means a backend added a dep with an `until` tensor that is not a node of the optimized graph
-    GGML_ASSERT(n_dep_nodes_added == n_dep_nodes);
-
     if (sched->n_copies > 1) {
         // add input copies as leafs so that they are allocated first
         for (int i = 0; i < sched->n_graph_inputs; i++) {

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

@@ -2936,13 +2936,12 @@ struct ggml_cplan ggml_graph_plan(
                         const int64_t ne10 = node->src[1]->ne[0]; // W
                         const int64_t ne11 = node->src[1]->ne[1]; // H
                         const int64_t ne12 = node->src[1]->ne[2]; // Channels In
-                        const int64_t ne13 = node->src[1]->ne[3]; // Batch
 
                         GGML_ASSERT(node->src[0]->type == GGML_TYPE_F16 || node->src[0]->type == GGML_TYPE_F32);
                         GGML_ASSERT(node->src[1]->type == GGML_TYPE_F32);
 
                         cur += ggml_type_size(node->src[0]->type) * ne00 * ne01 * ne02 * ne03;
-                        cur += ggml_type_size(node->src[0]->type) * ne10 * ne11 * ne12 * ne13;
+                        cur += ggml_type_size(node->src[0]->type) * ne10 * ne11 * ne12;
 
                     } break;
                 case GGML_OP_TOP_K:

ggml/src/ggml-cpu/ops.cpp

@@ -7267,21 +7267,18 @@ static void ggml_compute_forward_conv_transpose_2d_impl(
             }
         }
 
-        // permute source data (src1) from (Sw x Sh x Cin) to (Cin x Sw x Sh), for all batches
+        // permute source data (src1) from (Sw x Sh x Cin) to (Cin x Sw x Sh)
         {
             kernel_t * const wdata = (kernel_t *) params->wdata + nk;
-            for (int i13 = 0; i13 < ne13; i13++) {
-                kernel_t * const wdata_b = wdata + i13*ne10*ne11*ne12;
-                for (int i12 = 0; i12 < ne12; i12++) {
-                    for (int i11 = 0; i11 < ne11; i11++) {
-                        const float * const src = (float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11);
-                        kernel_t * dst_data = wdata_b + i11*ne10*ne12;
-                        for (int i10 = 0; i10 < ne10; i10++) {
-                            if constexpr (std::is_same_v<kernel_t, ggml_fp16_t>) {
-                                dst_data[i10*ne12 + i12] = GGML_CPU_FP32_TO_FP16(src[i10]);
-                            } else {
-                                dst_data[i10*ne12 + i12] = src[i10];
-                            }
+            for (int i12 = 0; i12 < ne12; i12++) {
+                for (int i11 = 0; i11 < ne11; i11++) {
+                    const float * const src = (float *)((char *) src1->data + i12*nb12 + i11*nb11);
+                    kernel_t * dst_data = wdata + i11*ne10*ne12;
+                    for (int i10 = 0; i10 < ne10; i10++) {
+                        if constexpr (std::is_same_v<kernel_t, ggml_fp16_t>) {
+                            dst_data[i10*ne12 + i12] = GGML_CPU_FP32_TO_FP16(src[i10]);
+                        } else {
+                            dst_data[i10*ne12 + i12] = src[i10];
                         }
                     }
                 }
@@ -7308,27 +7305,24 @@ static void ggml_compute_forward_conv_transpose_2d_impl(
     kernel_t * const wdata_src = wdata + nk;
 
     for (int i2 = ip0; i2 < ip1; i2++) { // Cout
+        float * dst_data = (float *)((char *) dst->data + i2*nb2);
         kernel_t * wdata_kernel = wdata + i2*ne01*ne00*ne03;
-        for (int i3 = 0; i3 < ne3; i3++) { // batch
-            float * dst_data = (float *)((char *) dst->data + i3*nb3 + i2*nb2);
-            kernel_t * wdata_src_b = wdata_src + i3*ne10*ne11*ne12;
-            for (int i11 = 0; i11 < ne11; i11++) {
-                for (int i10 = 0; i10 < ne10; i10++) {
-                    const int i1n = i11*ne10*ne12 + i10*ne12;
-                    for (int i01 = 0; i01 < ne01; i01++) {
-                        for (int i00 = 0; i00 < ne00; i00++) {
-                            float v = 0;
-                            if constexpr (std::is_same_v<kernel_t, ggml_fp16_t>) {
-                                ggml_vec_dot_f16(ne03, &v, 0,
-                                        wdata_src_b + i1n, 0,
-                                        wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1);
-                            } else {
-                                ggml_vec_dot_f32(ne03, &v, 0,
-                                        wdata_src_b + i1n, 0,
-                                        wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1);
-                            }
-                            dst_data[(i11*stride + i01)*ne0 + i10*stride + i00] += v;
+        for (int i11 = 0; i11 < ne11; i11++) {
+            for (int i10 = 0; i10 < ne10; i10++) {
+                const int i1n = i11*ne10*ne12 + i10*ne12;
+                for (int i01 = 0; i01 < ne01; i01++) {
+                    for (int i00 = 0; i00 < ne00; i00++) {
+                        float v = 0;
+                        if constexpr (std::is_same_v<kernel_t, ggml_fp16_t>) {
+                            ggml_vec_dot_f16(ne03, &v, 0,
+                                    wdata_src + i1n, 0,
+                                    wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1);
+                        } else {
+                            ggml_vec_dot_f32(ne03, &v, 0,
+                                    wdata_src + i1n, 0,
+                                    wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1);
                         }
+                        dst_data[(i11*stride + i01)*ne0 + i10*stride + i00] += v;
                     }
                 }
             }

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

@@ -4328,9 +4328,7 @@ static void ggml_backend_cuda_event_wait(ggml_backend_t backend, ggml_backend_ev
     }
 }
 
-static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph, ggml_backend_graph_optimize_params * params) {
-    GGML_UNUSED(params);
-
+static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph) {
     ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context;
 
 #ifdef USE_CUDA_GRAPH

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

@@ -69,15 +69,30 @@ using u32vec  = std::vector<uint32_t>;
 #define GGML_HEXAGON_FENCE_SLOT_SIZE       128
 
 struct ggml_hexagon_device_config {
-    int         physical_idx = 0;
-    int         virtual_idx  = 0;
-    int         domain_id    = 0;
-    std::string domain_name;
+    int physical_idx = 0;
+    int virtual_idx  = 0;
     std::string name;
 };
 
 static ggml_hexagon_device_config opt_device_configs[GGML_HEXAGON_MAX_SESSIONS];
 
+static int get_domain_id(int physical_idx) {
+    switch (physical_idx) {
+        case 0:  return 3;  // CDSP0 (all devices)
+        case 1:  return 4;  // CDSP1 (IQ9, IQ10)
+        case 2:  return 18; // CDSP2 (IQ10)
+        case 3:  return 19; // CDSP3 (IQ10)
+        default: return CDSP_DOMAIN_ID + physical_idx;
+    }
+}
+
+static std::string get_domain_name(int physical_idx) {
+    if (physical_idx == 0) {
+        return CDSP_DOMAIN_NAME;
+    }
+    return std::string("cdsp") + std::to_string(physical_idx);
+}
+
 static int    opt_arch    = 0; // autodetect
 static size_t opt_ndev    = 1;
 static size_t opt_nhvx    = 0; // use all
@@ -346,6 +361,7 @@ struct ggml_hexagon_session {
     uint32_t         session_id;
     uint32_t         domain_id;
     uint64_t         queue_id;
+    int              dev_id;
     int              phys_idx;
     int              virt_idx;
     bool             valid_session;
@@ -360,6 +376,9 @@ struct ggml_hexagon_session {
     std::unordered_map<int, std::unique_ptr<ggml_hexagon_shared_buffer>> cloned_buffers;
     std::unordered_set<ggml_hexagon_session *>                           sync_peers;
 
+    ggml_backend_buffer_type buffer_type        = {};
+    ggml_backend_buffer_type host_buffer_type   = {};
+
     uint32_t n_threads   = 0;
     uint32_t n_hvx       = 0;
     uint32_t n_hmx       = 0;
@@ -373,12 +392,12 @@ struct ggml_hexagon_session {
 
     mutable std::unordered_set<const ggml_tensor *> needs_repack;
 
-    ggml_hexagon_session(const ggml_hexagon_device_config & config, ggml_backend_dev_t dev = nullptr) noexcept(false);
+    ggml_hexagon_session(int dev_id, ggml_backend_dev_t dev) noexcept(false);
     ~ggml_hexagon_session() noexcept(true);
 
     const char* c_name() const { return name.c_str(); }
 
-    void allocate(const ggml_hexagon_device_config & config) noexcept(false);
+    void allocate(int dev_id) noexcept(false);
     void release() noexcept(true);
 
     void enqueue_op(const htp_opnode & node);
@@ -411,38 +430,14 @@ struct ggml_hexagon_session {
 
 // ** backend buffers
 
-struct ggml_backend_hexagon_device_context {
-    int                        dev_id;
-    ggml_hexagon_device_config config;
-    ggml_backend_dev_t         dev = nullptr;
-    size_t                     max_bufsize = 0;
-
-    ggml_backend_buffer_type buffer_type      = {};
-    ggml_backend_buffer_type host_buffer_type = {};
-
-    std::unique_ptr<ggml_hexagon_session> sess;
-
-    ggml_backend_hexagon_device_context(int dev_id, const ggml_hexagon_device_config & config, ggml_backend_dev_t dev);
-    ~ggml_backend_hexagon_device_context();
-
-    const char * c_name() const { return config.name.c_str(); }
-
-    ggml_hexagon_session * session() {
-        if (!sess) {
-            sess = std::make_unique<ggml_hexagon_session>(config, dev);
-        }
-        return sess.get();
-    }
-};
-
 struct ggml_backend_hexagon_buffer_type_context {
-    ggml_backend_hexagon_buffer_type_context(const std::string & name, ggml_backend_hexagon_device_context * dev_ctx) {
-        this->dev_ctx = dev_ctx;
-        this->name    = name;
+    ggml_backend_hexagon_buffer_type_context(const std::string & name, ggml_hexagon_session * sess) {
+        this->sess = sess;
+        this->name = name;
     }
 
-    ggml_backend_hexagon_device_context * dev_ctx;
-    std::string                           name;
+    ggml_hexagon_session * sess;
+    std::string            name;
 };
 
 struct ggml_hexagon_rpcmem_block {
@@ -581,8 +576,7 @@ struct ggml_hexagon_shared_buffer {
 };
 
 static ggml_hexagon_session * ggml_backend_hexagon_buffer_get_sess(ggml_backend_buffer_t buffer) {
-    auto sbuf = static_cast<ggml_hexagon_shared_buffer *>(buffer->context);
-    return sbuf->sess;
+    return static_cast<ggml_backend_hexagon_buffer_type_context *>(buffer->buft->context)->sess;
 }
 
 static void ggml_backend_hexagon_buffer_free_buffer(ggml_backend_buffer_t buffer) {
@@ -1500,26 +1494,24 @@ static const char * ggml_backend_hexagon_buffer_type_name(ggml_backend_buffer_ty
 
 static ggml_backend_buffer_t ggml_backend_hexagon_buffer_type_alloc_buffer(
             ggml_backend_buffer_type_t buffer_type, size_t size) {
-    auto dev_ctx = static_cast<ggml_backend_hexagon_buffer_type_context *>(buffer_type->context)->dev_ctx;
-    auto sess    = dev_ctx->session();
+    auto sess = static_cast<ggml_backend_hexagon_buffer_type_context *>(buffer_type->context)->sess;
     try {
         ggml_hexagon_shared_buffer * sbuf = new ggml_hexagon_shared_buffer(sess, size, false, GGML_HEXAGON_FENCE_BUFFER_SIZE);
         return ggml_backend_buffer_init(buffer_type, ggml_backend_hexagon_buffer_interface, sbuf, size);
     } catch (const std::exception & exc) {
-        GGML_LOG_ERROR("ggml-hex: %s failed to allocate device buffer context: %s\n", dev_ctx->c_name(), exc.what());
+        GGML_LOG_ERROR("ggml-hex: %s failed to allocate device buffer context: %s\n", sess->c_name(), exc.what());
         return nullptr;
     }
 }
 
 static ggml_backend_buffer_t ggml_backend_hexagon_host_buffer_type_alloc_buffer(
             ggml_backend_buffer_type_t buffer_type, size_t size) {
-    auto dev_ctx = static_cast<ggml_backend_hexagon_buffer_type_context *>(buffer_type->context)->dev_ctx;
-    auto sess    = dev_ctx->session();
+    auto sess = static_cast<ggml_backend_hexagon_buffer_type_context *>(buffer_type->context)->sess;
     try {
         ggml_hexagon_shared_buffer * sbuf = new ggml_hexagon_shared_buffer(sess, size, false, GGML_HEXAGON_FENCE_BUFFER_SIZE);
         return ggml_backend_buffer_init(buffer_type, ggml_backend_hexagon_host_buffer_interface, sbuf, size);
     } catch (const std::exception & exc) {
-        GGML_LOG_ERROR("ggml-hex: %s failed to allocate host buffer context: %s\n", dev_ctx->c_name(), exc.what());
+        GGML_LOG_ERROR("ggml-hex: %s failed to allocate host buffer context: %s\n", sess->c_name(), exc.what());
         return nullptr;
     }
 }
@@ -1544,7 +1536,7 @@ static size_t ggml_backend_hexagon_buffer_type_get_alloc_size(ggml_backend_buffe
 
 static size_t ggml_backend_hexagon_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) {
     auto * context = static_cast<ggml_backend_hexagon_buffer_type_context *>(buft->context);
-    return context->dev_ctx->max_bufsize;
+    return context->sess->max_bufsize;
 }
 
 static bool ggml_backend_hexagon_buffer_type_is_host(ggml_backend_buffer_type_t buft) {
@@ -1575,22 +1567,6 @@ static ggml_backend_buffer_type_i ggml_backend_hexagon_host_buffer_type_interfac
     /* .is_host          = */ ggml_backend_hexagon_host_buffer_type_is_host,
 };
 
-ggml_backend_hexagon_device_context::ggml_backend_hexagon_device_context(int dev_id, const ggml_hexagon_device_config & config, ggml_backend_dev_t dev)
-    : dev_id(dev_id), config(config), dev(dev), max_bufsize(opt_mbuf) {
-    buffer_type.device  = dev;
-    buffer_type.iface   = ggml_backend_hexagon_buffer_type_interface;
-    buffer_type.context = new ggml_backend_hexagon_buffer_type_context(config.name, this);
-
-    host_buffer_type.device  = dev;
-    host_buffer_type.iface   = ggml_backend_hexagon_host_buffer_type_interface;
-    host_buffer_type.context = new ggml_backend_hexagon_buffer_type_context(config.name + "-HOST", this);
-}
-
-ggml_backend_hexagon_device_context::~ggml_backend_hexagon_device_context() {
-    delete static_cast<ggml_backend_hexagon_buffer_type_context *>(buffer_type.context);
-    delete static_cast<ggml_backend_hexagon_buffer_type_context *>(host_buffer_type.context);
-}
-
 static bool ggml_backend_buffer_is_hexagon(const struct ggml_backend_buffer * b) {
     return b->buft->iface.get_alignment == ggml_backend_hexagon_buffer_type_get_alignment;
 }
@@ -2835,7 +2811,8 @@ static size_t ggml_hexagon_measure_max_vmem(ggml_hexagon_session *sess) {
     return vmem - step; // backoff to account for overhead from internal mappings
 }
 
-void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) noexcept(false) {
+void ggml_hexagon_session::allocate(int dev_id) noexcept(false) {
+    const auto & config = opt_device_configs[dev_id];
     int phys_idx = config.physical_idx;
     int virt_idx = config.virtual_idx;
 
@@ -2846,31 +2823,21 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
 
     this->phys_idx   = phys_idx;
     this->virt_idx   = virt_idx;
-    this->domain_id  = config.domain_id;
+    this->domain_id  = get_domain_id(phys_idx);
     this->session_id = 0;
+    this->dev_id     = dev_id;
     this->name       = config.name;
     this->op_pending = 0;
 
     GGML_LOG_DEBUG("ggml-hex: %s allocating new session\n", this->name.c_str());
 
-    if (config.domain_id < 0 || config.domain_name.empty()) {
-        GGML_LOG_ERROR("ggml-hex: %s: invalid physical CDSP core %d\n", config.name.c_str(), config.physical_idx);
-        throw std::runtime_error("ggml-hex: invalid physical CDSP core");
+    domain * my_domain = htpdrv_get_domain(this->domain_id);
+    if (my_domain == NULL) {
+        GGML_LOG_ERROR("ggml-hex: unable to get domain struct for CDSP (domain_id %d)\n", this->domain_id);
+        throw std::runtime_error("ggml-hex: failed to get CDSP domain (see log for details)");
     }
 
-    const std::string & dom_name = config.domain_name;
-
-    // Enable Unsigned PD for all domains
-    {
-        struct remote_rpc_control_unsigned_module u;
-        u.domain = -1;
-        u.enable = 1;
-        int err  = remote_session_control(DSPRPC_CONTROL_UNSIGNED_MODULE, (void *) &u, sizeof(u));
-        if (err != AEE_SUCCESS) {
-            GGML_LOG_ERROR("ggml-hex: %s failed to enable unsigned PD : error 0x%x\n", this->c_name(), err);
-            throw std::runtime_error("ggml-hex: remote_session_control(unsign) failed (see log for details)");
-        }
-    }
+    std::string dom_name = get_domain_name(phys_idx);
 
     // Create new session if virtual_idx > 0
     if (virt_idx > 0) {
@@ -2882,8 +2849,7 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
 
         int err = remote_session_control(FASTRPC_RESERVE_NEW_SESSION, (void *) &n, sizeof(n));
         if (err != AEE_SUCCESS) {
-            GGML_LOG_ERROR("ggml-hex: %s failed to reserve new session (physical %d, virtual %d) : error 0x%x\n",
-                           this->c_name(), phys_idx, virt_idx, err);
+            GGML_LOG_ERROR("ggml-hex: failed to reserve new session %d (physical %d, virtual %d) : error 0x%x\n", dev_id, phys_idx, virt_idx, err);
             throw std::runtime_error("ggml-hex: remote_session_control(new-sess) failed (see log for details)");
         }
 
@@ -2891,21 +2857,10 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
         this->session_id    = n.session_id;
         this->domain_id     = n.effective_domain_id;
         this->valid_session = true;
-    } else {
-        struct remote_rpc_effective_domain_id eff = {};
-        eff.domain_name     = const_cast<char *>(dom_name.c_str());
-        eff.domain_name_len = dom_name.size();
-        eff.session_id      = 0;
-
-        int err = remote_session_control(FASTRPC_GET_EFFECTIVE_DOMAIN_ID, (void *) &eff, sizeof(eff));
-        if (err == AEE_SUCCESS) {
-            this->domain_id = eff.effective_domain_id;
-        } else {
-            GGML_LOG_DEBUG("ggml-hex: %s FASTRPC_GET_EFFECTIVE_DOMAIN_ID returned 0x%x, using domain_id %d\n",
-                           this->name.c_str(), err, this->domain_id);
-        }
     }
 
+    // Get session URI
+
     char session_uri[256];
     {
         char htp_uri[256];
@@ -2922,18 +2877,31 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
 
         int err = remote_session_control(FASTRPC_GET_URI, (void *) &u, sizeof(u));
         if (err != AEE_SUCCESS) {
-            snprintf(session_uri, sizeof(session_uri), "%s&_dom=%s&_session=%u",
-                     htp_uri, dom_name.c_str(), this->session_id);
+            // fallback to single session uris
+            int htp_URI_domain_len = strlen(htp_uri) + MAX_DOMAIN_NAMELEN;
 
-            GGML_LOG_WARN("ggml-hex: %s failed to get URI (physical %d, virtual %d) : error 0x%x. Falling back to single session URI: %s\n",
-                          this->c_name(), phys_idx, virt_idx, err, session_uri);
+            snprintf(session_uri, htp_URI_domain_len, "%s%s", htp_uri, my_domain->uri);
+
+            GGML_LOG_WARN("ggml-hex: failed to get URI for session %d (physical %d, virtual %d) : error 0x%x. Falling back to single session URI: %s\n", dev_id, phys_idx, virt_idx, err, session_uri);
+        }
+    }
+
+    // Enable Unsigned PD
+    {
+        struct remote_rpc_control_unsigned_module u;
+        u.domain = this->domain_id;
+        u.enable = 1;
+        int err  = remote_session_control(DSPRPC_CONTROL_UNSIGNED_MODULE, (void *) &u, sizeof(u));
+        if (err != AEE_SUCCESS) {
+            GGML_LOG_ERROR("ggml-hex: failed to enable unsigned PD for session %d : error 0x%x\n", dev_id, err);
+            throw std::runtime_error("ggml-hex: remote_session_control(unsign) failed (see log for details)");
         }
     }
 
     // Open session
     int err = htp_iface_open(session_uri, &this->handle);
     if (err != AEE_SUCCESS) {
-        GGML_LOG_ERROR("ggml-hex: %s failed to open session : error 0x%x\n", this->c_name(), err);
+        GGML_LOG_ERROR("ggml-hex: failed to open session %d : error 0x%x\n", dev_id, err);
         throw std::runtime_error("ggml-hex: failed to open session (see log for details)");
     }
 
@@ -3023,7 +2991,7 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
     this->op_batch = new ggml_hexagon_opbatch(this, opt_opbatch, this->max_vmem);
 
     // Start dspqueue/opbatch processing
-    err = htp_iface_start(this->handle, this->session_id, this->queue_id, opt_nhvx, opt_nhmx, this->max_vmem);
+    err = htp_iface_start(this->handle, dev_id, this->queue_id, opt_nhvx, opt_nhmx, this->max_vmem);
     if (err != 0) {
         GGML_LOG_ERROR("ggml-hex: %s failed to start session: 0x%08x\n", this->c_name(), (unsigned) err);
         throw std::runtime_error("ggml-hex: iface start failed (see log for details)");
@@ -3086,23 +3054,33 @@ void ggml_hexagon_session::release() noexcept(true) {
     this->cloned_buffers.clear();
 }
 
-ggml_hexagon_session::ggml_hexagon_session(const ggml_hexagon_device_config & config, ggml_backend_dev_t dev) noexcept(false) {
+ggml_hexagon_session::ggml_hexagon_session(int dev_id, ggml_backend_dev_t dev) noexcept(false) {
+    buffer_type.device      = dev;
+    host_buffer_type.device = dev;
+
     op_batch = nullptr;
     op_queue = nullptr;
     fence_seq = ((uintptr_t)this) & 0xFFFF;
 
     try {
-        allocate(config);
+        allocate(dev_id);
+
+        buffer_type.iface   = ggml_backend_hexagon_buffer_type_interface;
+        buffer_type.context = new ggml_backend_hexagon_buffer_type_context(this->name, this);
+
+        host_buffer_type.iface   = ggml_backend_hexagon_host_buffer_type_interface;
+        host_buffer_type.context = new ggml_backend_hexagon_buffer_type_context(this->name + "-HOST", this);
     } catch (const std::exception & exc) {
         release();
         throw;
     }
-
-    GGML_UNUSED(dev);
 }
 
 ggml_hexagon_session::~ggml_hexagon_session() noexcept(true) {
     release();
+
+    delete static_cast<ggml_backend_hexagon_buffer_type_context *>(buffer_type.context);
+    delete static_cast<ggml_backend_hexagon_buffer_type_context *>(host_buffer_type.context);
 }
 
 // ** backend interface
@@ -3979,13 +3957,11 @@ static void ggml_hexagon_precompute_fused_mmnx_params(
 }
 
 static bool ggml_hexagon_tensor_is_host(const struct ggml_hexagon_session * sess, const struct ggml_tensor * t) {
-    return t && t->buffer && ggml_backend_buft_is_host(t->buffer->buft);
-    GGML_UNUSED(sess);
+    return t && t->buffer && t->buffer->buft == &sess->host_buffer_type;
 }
 
 static bool ggml_hexagon_tensor_is_non_host(const struct ggml_hexagon_session * sess, const struct ggml_tensor * t) {
-    return t && t->buffer && !ggml_backend_buft_is_host(t->buffer->buft);
-    GGML_UNUSED(sess);
+    return t && t->buffer && t->buffer->buft != &sess->host_buffer_type;
 }
 
 static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * sess, const struct ggml_tensor * dst) {
@@ -5008,9 +4984,7 @@ static std::vector<int> ggml_hexagon_graph_optimize_reorder(const std::vector<ht
     return res;
 }
 
-static void ggml_backend_hexagon_graph_optimize(ggml_backend_t backend, ggml_cgraph * gf, ggml_backend_graph_optimize_params * params) {
-    GGML_UNUSED(params);
-
+static void ggml_backend_hexagon_graph_optimize(ggml_backend_t backend, ggml_cgraph * gf) {
     const int n = gf->n_nodes;
 
     constexpr int MAX_FUSE = 16;
@@ -5293,8 +5267,7 @@ bool ggml_backend_is_hexagon(ggml_backend_t backend) {
 // device interface
 
 static ggml_backend_t ggml_backend_hexagon_device_init(ggml_backend_dev_t dev, const char * params) {
-    auto dev_ctx = static_cast<ggml_backend_hexagon_device_context *>(dev->context);
-    auto sess    = dev_ctx->session();
+    auto sess = static_cast<ggml_hexagon_session *>(dev->context);
 
     return new ggml_backend{
         /* .guid      = */ ggml_backend_hexagon_guid(),
@@ -5307,8 +5280,8 @@ static ggml_backend_t ggml_backend_hexagon_device_init(ggml_backend_dev_t dev, c
 }
 
 static const char * ggml_backend_hexagon_device_get_name(ggml_backend_dev_t dev) {
-    auto dev_ctx = static_cast<ggml_backend_hexagon_device_context *>(dev->context);
-    return dev_ctx->c_name();
+    auto sess = static_cast<ggml_hexagon_session *>(dev->context);
+    return sess->c_name();
 
     GGML_UNUSED(dev);
 }
@@ -5346,16 +5319,16 @@ static void ggml_backend_hexagon_device_get_props(ggml_backend_dev_t dev, struct
 }
 
 static ggml_backend_buffer_type_t ggml_backend_hexagon_device_get_buffer_type(ggml_backend_dev_t dev) {
-    auto dev_ctx = static_cast<ggml_backend_hexagon_device_context *>(dev->context);
-    return &dev_ctx->buffer_type;
+    auto sess = static_cast<ggml_hexagon_session *>(dev->context);
+    return &sess->buffer_type;
 }
 
 static ggml_backend_buffer_type_t ggml_backend_hexagon_device_get_host_buffer_type(ggml_backend_dev_t dev) {
     if (!opt_hostbuf) {
         return NULL;
     }
-    auto dev_ctx = static_cast<ggml_backend_hexagon_device_context *>(dev->context);
-    return &dev_ctx->host_buffer_type;
+    auto sess = static_cast<ggml_hexagon_session *>(dev->context);
+    return &sess->host_buffer_type;
 }
 
 static bool ggml_hexagon_supported_cpy(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -5446,8 +5419,7 @@ static bool ggml_hexagon_supported_fill(const struct ggml_hexagon_session * sess
 }
 
 static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) {
-    auto dev_ctx = static_cast<ggml_backend_hexagon_device_context *>(dev->context);
-    auto sess    = dev_ctx->session();
+    auto sess = static_cast<ggml_hexagon_session *>(dev->context);
 
     // reject ops that match the filter
     if (opt_opfilter && std::regex_match(ggml_op_desc(op), *opt_opfilter)) {
@@ -5519,7 +5491,6 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
                     supp = ggml_hexagon_supported_unary(sess, op);
                     break;
                 default:
-                    supp = false;
                     break;
             }
             break;
@@ -5532,7 +5503,6 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
                     supp = ggml_hexagon_supported_activations(sess, op);
                     break;
                 default:
-                    supp = false;
                     break;
             }
             break;
@@ -5618,17 +5588,17 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
 }
 
 static bool ggml_backend_hexagon_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) {
-    auto dev_ctx = static_cast<ggml_backend_hexagon_device_context *>(dev->context);
+    auto sess = static_cast<ggml_hexagon_session *>(dev->context);
 
     // Technically we can clone hexagon buffers from any session but for some reason the output is garbled with layer-split,
     // tensor-split works correctly, so it needs mode debugging and investigation. For now accept only our own buffers.
 #if 0
     bool supp = (buft->iface.get_alignment == ggml_backend_hexagon_buffer_type_get_alignment);
 #else
-    bool supp = (buft == &dev_ctx->host_buffer_type) || (buft == &dev_ctx->buffer_type);
+    bool supp = (buft == &sess->host_buffer_type) || (buft == &sess->buffer_type);
 #endif
 
-    HEX_VERBOSE("ggml-hex: %s device-supports-buft %s %s\n", dev_ctx->c_name(), ggml_backend_buft_name(buft), supp ? "yes" : "no");
+    HEX_VERBOSE("ggml-hex: %s device-supports-buft %s %s\n", sess->name.c_str(), ggml_backend_buft_name(buft), supp ? "yes" : "no");
     return supp;
 }
 
@@ -5657,11 +5627,16 @@ ggml_hexagon_registry::ggml_hexagon_registry(ggml_backend_reg_t reg) {
 
     GGML_LOG_INFO("ggml-hex: Hexagon Arch version v%d\n", opt_arch);
 
-    // Create devices
+    // Create devices / sessions
     for (size_t i = 0; i < opt_ndev; i++) {
-        devices[i].iface   = ggml_backend_hexagon_device_i;
-        devices[i].reg     = reg;
-        devices[i].context = new ggml_backend_hexagon_device_context(i, opt_device_configs[i], &devices[i]);
+        devices[i].iface = ggml_backend_hexagon_device_i;
+        devices[i].reg   = reg;
+        try {
+            devices[i].context = new ggml_hexagon_session(i, &devices[i]);
+        } catch (const std::exception & exc) {
+            GGML_LOG_ERROR("ggml-hex: failed to create device/session %zu\n", i);
+            devices[i].context = nullptr;
+        }
     }
 
 }
@@ -5669,10 +5644,10 @@ ggml_hexagon_registry::ggml_hexagon_registry(ggml_backend_reg_t reg) {
 ggml_hexagon_registry::~ggml_hexagon_registry() {
     GGML_LOG_INFO("ggml-hex: releasing registry\n");
 
-    // Release devices
+    // Release devices / sessions
     for (size_t i = 0; i < opt_ndev; i++) {
-        auto dev_ctx = static_cast<ggml_backend_hexagon_device_context *>(devices[i].context);
-        delete dev_ctx;
+        auto sess = static_cast<ggml_hexagon_session *>(devices[i].context);
+        delete sess;
     }
 }
 
@@ -5841,85 +5816,6 @@ template<typename T, int BASE=10> std::string vec_to_str(std::vector<T> v) {
     return str;
 }
 
-// Enumerate NPU (aka CDSP) domains via FASTRPC_GET_DOMAINS if supported,
-// and populate domain_id and domain_name for all configured devices.
-static void ggml_hexagon_discover_devices() {
-    std::unordered_map<int, fastrpc_domain> cdsp_map;
-    bool discovery_supported = false;
-
-    system_req_payload domain_info = {};
-    domain_info.id              = FASTRPC_GET_DOMAINS;
-    domain_info.sys.domains     = nullptr;
-    domain_info.sys.max_domains = 0;
-    domain_info.sys.flags       = DOMAINS_LIST_FLAGS_SET_TYPE(0, FASTRPC_NSP);
-
-    int err = remote_system_request(&domain_info);
-    if (err == AEE_SUCCESS && domain_info.sys.num_domains > 0) {
-        std::vector<fastrpc_domain> domains(domain_info.sys.num_domains);
-        domain_info.sys.domains     = domains.data();
-        domain_info.sys.max_domains = (int) domains.size();
-
-        err = remote_system_request(&domain_info);
-        if (err == AEE_SUCCESS) {
-            discovery_supported = true;
-            const int n_domains = std::min(domain_info.sys.num_domains, (int) domains.size());
-            for (int i = 0; i < n_domains; i++) {
-                GGML_LOG_INFO("ggml-hex: FASTRPC_GET_DOMAINS[%d]: type %d id %d name '%s' status %d instance-id %d\n",
-                              i, (int) domains[i].type, domains[i].id, domains[i].name, domains[i].status, domains[i].instance_id);
-                if (domains[i].type != FASTRPC_NSP) {
-                    GGML_LOG_DEBUG("ggml-hex:   skipping non-CDSP domain (type=%d)\n", (int) domains[i].type);
-                    continue;
-                }
-                if (!domains[i].status) {
-                    GGML_LOG_WARN("ggml-hex:   skipping CDSP domain id=%d (status=down)\n", domains[i].id);
-                    continue;
-                }
-                cdsp_map[domains[i].instance_id] = domains[i];
-                GGML_LOG_INFO("ggml-hex: using CDSP domain: instance-id %d id %d name '%s'\n",
-                              domains[i].instance_id, domains[i].id, domains[i].name);
-            }
-        } else {
-            GGML_LOG_WARN("ggml-hex: FASTRPC_GET_DOMAINS fetch failed (0x%x), using static CDSP domains\n", (unsigned) err);
-        }
-    } else if (err != AEE_SUCCESS) {
-        GGML_LOG_DEBUG("ggml-hex: FASTRPC_GET_DOMAINS query failed (0x%x), using static CDSP domains\n", (unsigned) err);
-    }
-
-    // Populate domain IDs and names for all configured devices
-    for (size_t i = 0; i < opt_ndev; i++) {
-        auto & cfg = opt_device_configs[i];
-        if (discovery_supported) {
-            auto it = cdsp_map.find(cfg.physical_idx);
-            if (it != cdsp_map.end()) {
-                cfg.domain_id   = it->second.id;
-                cfg.domain_name = it->second.name;
-            } else {
-                GGML_LOG_ERROR("ggml-hex: physical CDSP core %d not found on device (%zu CDSP core(s) available)\n",
-                               cfg.physical_idx, cdsp_map.size());
-                cfg.domain_id   = -1;
-                cfg.domain_name = "";
-            }
-        } else {
-            switch (cfg.physical_idx) {
-                case 0:
-                    cfg.domain_id   = 3;
-                    cfg.domain_name = CDSP_DOMAIN_NAME;
-                    break;
-                case 1:
-                    cfg.domain_id   = 4;
-                    cfg.domain_name = "cdsp1";
-                    break;
-                default:
-                    GGML_LOG_ERROR("ggml-hex: physical CDSP core %d not supported without dynamic discovery\n",
-                                   cfg.physical_idx);
-                    cfg.domain_id   = -1;
-                    cfg.domain_name = "";
-                    break;
-            }
-        }
-    }
-}
-
 static void ggml_hexagon_init(ggml_backend_reg * reg) {
     // Basic sanity checks to make sure definitions match
     static_assert((unsigned int) HTP_TYPE_Q4_0 == (unsigned int) GGML_TYPE_Q4_0,
@@ -6085,9 +5981,6 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
     }
 #endif
 
-    // Resolve domain info for all configured devices
-    ggml_hexagon_discover_devices();
-
     if (str_profile) {
         opt_pmu_evt = [&]() -> std::vector<uint32_t> {
             auto v  = str_to_vec<uint32_t>(str_profile);

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

@@ -73,7 +73,6 @@ typedef int (*remote_handle64_close_pfn_t)(remote_handle h);
 typedef int (*remote_handle_control_pfn_t)(uint32_t req, void* data, uint32_t datalen);
 typedef int (*remote_handle64_control_pfn_t)(remote_handle64 h, uint32_t req, void* data, uint32_t datalen);
 typedef int (*remote_session_control_pfn_t)(uint32_t req, void *data, uint32_t datalen);
-typedef int (*remote_system_request_pfn_t)(system_req_payload * req);
 
 //
 // Driver API pfns
@@ -100,7 +99,6 @@ remote_handle64_close_pfn_t   remote_handle64_close_pfn   = nullptr;
 remote_handle_control_pfn_t   remote_handle_control_pfn   = nullptr;
 remote_handle64_control_pfn_t remote_handle64_control_pfn = nullptr;
 remote_session_control_pfn_t  remote_session_control_pfn  = nullptr;
-remote_system_request_pfn_t   remote_system_request_pfn   = nullptr;
 
 //
 // Driver API
@@ -208,13 +206,6 @@ HTPDRV_API int remote_session_control(uint32_t req, void * data, uint32_t datale
     return remote_session_control_pfn(req, data, datalen);
 }
 
-HTPDRV_API int remote_system_request(system_req_payload * req) {
-    if (!remote_system_request_pfn) {
-        return AEE_EUNSUPPORTEDAPI;
-    }
-    return remote_system_request_pfn(req);
-}
-
 #ifdef _WIN32
 
 static std::string wstr_to_str(std::wstring_view wstr) {
@@ -376,7 +367,6 @@ int htpdrv_init() {
     dlsym(handle.get(), remote_handle64_control_pfn_t, remote_handle64_control_pfn, remote_handle64_control, false);
     dlsym(handle.get(), remote_session_control_pfn_t, remote_session_control_pfn, remote_session_control, false);
     dlsym(handle.get(), remote_handle64_close_pfn_t, remote_handle64_close_pfn, remote_handle64_close, false);
-    dlsym(handle.get(), remote_system_request_pfn_t, remote_system_request_pfn, remote_system_request, true);
 
     lib_cdsp_rpc_handle = std::move(handle);
     initialized         = true;

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

@@ -116,8 +116,6 @@ HTPDRV_API domain * htpdrv_get_domain(int domain_id);
  */
 HTPDRV_API int htpdrv_get_arch(int domain, int * arch);
 
-HTPDRV_API int remote_system_request(system_req_payload * req);
-
 #ifdef __cplusplus
 }
 #endif

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

@@ -1,27 +1,10 @@
 #include "ggml-metal-common.h"
 
-#include "ggml.h"
 #include "ggml-impl.h"
 #include "ggml-backend-impl.h"
 
 #include <vector>
 
-bool ggml_metal_op_mul_mat_use_mm(const struct ggml_tensor * op, bool has_simdgroup_mm) {
-    const int64_t ne00 = op->src[0]->ne[0];
-    const int64_t ne11 = op->src[1]->ne[1];
-
-    return !ggml_is_transposed(op->src[0]) &&
-           !ggml_is_transposed(op->src[1]) &&
-           has_simdgroup_mm && ne00 >= 64 && ne11 > 8;
-}
-
-bool ggml_metal_op_mul_mat_id_use_mm(const struct ggml_tensor * op, bool has_simdgroup_mm) {
-    const int64_t ne00 = op->src[0]->ne[0];
-    const int64_t ne21 = op->src[2]->ne[1];
-
-    return has_simdgroup_mm && ne00 >= 64 && ne21 >= 32;
-}
-
 // represents a memory range (i.e. an interval from a starting address p0 to an ending address p1 in a given buffer pb)
 // the type indicates whether it is a source range (i.e. ops read data from it) or a destination range (i.e. ops write data to it)
 struct ggml_mem_range {

ggml/src/ggml-metal/ggml-metal-common.h

@@ -47,10 +47,6 @@ bool ggml_mem_ranges_check(ggml_mem_ranges_t mrs, const struct ggml_tensor * ten
 //       if it proves to work well, we can start using it for other backends in the future
 void ggml_graph_optimize(struct ggml_cgraph * gf);
 
-// mat-mat vs mat-vec dispatch; used by both supports_op and ggml_metal_op_mul_mat*
-bool ggml_metal_op_mul_mat_use_mm   (const struct ggml_tensor * op, bool has_simdgroup_mm);
-bool ggml_metal_op_mul_mat_id_use_mm(const struct ggml_tensor * op, bool has_simdgroup_mm);
-
 #ifdef __cplusplus
 }
 #endif

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

@@ -593,7 +593,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_me
     // - sgptg floats for shared_x_dt (nsg)
     // - sgptg floats for shared_dA (nsg)
     // Total: nsg * (32 + 2) floats
-    res.smem = GGML_PAD((32 + 2)*sizeof(float)*nsg, 16);
+    res.smem = (32 + 2)*sizeof(float)*nsg;
 
     return res;
 }
@@ -1029,7 +1029,6 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0(g
     }
 
     res.smem = (size_t) ne02*ne20*sizeof(uint16_t);
-    res.smem = GGML_PAD(res.smem, 16);
 
     return res;
 }

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

@@ -3,7 +3,6 @@
 #import "ggml-impl.h"
 #import "ggml-backend-impl.h"
 #import "ggml-metal-impl.h"
-#import "ggml-metal-common.h"
 
 #include <Foundation/Foundation.h>
 
@@ -789,10 +788,6 @@ void ggml_metal_encoder_debug_group_pop (ggml_metal_encoder_t encoder) {
 }
 
 void ggml_metal_encoder_set_pipeline(ggml_metal_encoder_t encoder, struct ggml_metal_pipeline_with_params pipeline) {
-    if (!pipeline.pipeline) {
-        GGML_ABORT("%s: nil Metal pipeline (missing kernel; see compile_pipeline log above)\n", __func__);
-    }
-
     [encoder->obj setComputePipelineState:pipeline.pipeline->obj];
 }
 
@@ -805,10 +800,7 @@ void ggml_metal_encoder_set_buffer(ggml_metal_encoder_t encoder, struct ggml_met
 }
 
 void ggml_metal_encoder_set_threadgroup_memory_size(ggml_metal_encoder_t encoder, size_t size, int idx) {
-    // ref: https://developer.apple.com/documentation/metal/mtlcomputecommandencoder/setthreadgroupmemorylength(_:index:)
-    GGML_ASSERT(size % 16 == 0);
-
-    [encoder->obj setThreadgroupMemoryLength:size atIndex:idx];
+    [encoder->obj setThreadgroupMemoryLength:GGML_PAD(size, 16) atIndex:idx];
 }
 
 void ggml_metal_encoder_dispatch_threadgroups(ggml_metal_encoder_t encoder, int tg0, int tg1, int tg2, int tptg0, int tptg1, int tptg2) {
@@ -1415,30 +1407,6 @@ void ggml_metal_device_get_memory(ggml_metal_device_t dev, size_t * free, size_t
     }
 }
 
-static bool ggml_metal_supports_mul_mat_op(
-        bool has_simdgroup_reduction,
-        const struct ggml_tensor * op,
-        bool src0_f16_has_mv,
-        bool mm_path) {
-    if (!has_simdgroup_reduction || op->src[0]->type == GGML_TYPE_NVFP4) {
-        return false;
-    }
-
-    if (op->src[1]->type != GGML_TYPE_F16) {
-        return true;
-    }
-
-    if (op->src[0]->type == GGML_TYPE_BF16) {
-        return false;
-    }
-
-    if (src0_f16_has_mv && op->src[0]->type == GGML_TYPE_F16) {
-        return true;
-    }
-
-    return mm_path;
-}
-
 bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_tensor * op) {
     const bool has_simdgroup_mm        = dev->props.has_simdgroup_mm;
     const bool has_simdgroup_reduction = dev->props.has_simdgroup_reduction;
@@ -1742,15 +1710,9 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
         case GGML_OP_GATED_DELTA_NET:
             return has_simdgroup_reduction && op->src[2]->ne[0] % 32 == 0;
         case GGML_OP_SOLVE_TRI:
-            return has_simdgroup_reduction && op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_MUL_MAT:
-            return ggml_metal_supports_mul_mat_op(
-                    has_simdgroup_reduction, op, true,
-                    ggml_metal_op_mul_mat_use_mm(op, has_simdgroup_mm));
         case GGML_OP_MUL_MAT_ID:
-            return ggml_metal_supports_mul_mat_op(
-                    has_simdgroup_reduction, op, false,
-                    ggml_metal_op_mul_mat_id_use_mm(op, has_simdgroup_mm));
+            return has_simdgroup_reduction && op->src[0]->type != GGML_TYPE_NVFP4;
         case GGML_OP_SET:
         case GGML_OP_CPY:
         case GGML_OP_DUP:

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

@@ -660,7 +660,6 @@ typedef struct {
     uint64_t nb0;
     uint64_t nb1;
     uint64_t nb2;
-    uint64_t nb3;
 } ggml_metal_kargs_conv_transpose_2d;
 
 typedef struct {

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

@@ -948,7 +948,7 @@ int ggml_metal_op_sum(ggml_metal_op_t ctx, int idx) {
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         2);
 
-    ggml_metal_encoder_set_threadgroup_memory_size(enc, GGML_PAD(nsg * sizeof(float), 16), 0);
+    ggml_metal_encoder_set_threadgroup_memory_size(enc, nsg * sizeof(float), 0);
 
     ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, nth, 1, 1);
 
@@ -2362,6 +2362,10 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) {
     const int16_t r2 = ne12/ne02;
     const int16_t r3 = ne13/ne03;
 
+    // find the break-even point where the matrix-matrix kernel becomes more efficient compared
+    // to the matrix-vector kernel
+    const int ne11_mm_min = 8;
+
     // first try to use small-batch mat-mv kernels
     // these should be efficient for BS [2, ~8]
     if (op->src[1]->type == GGML_TYPE_F32 && (ne00%128 == 0) &&
@@ -2464,7 +2468,12 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) {
         ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         3);
 
         ggml_metal_encoder_dispatch_threadgroups(enc, ((ne01 + r0ptg - 1)/r0ptg), ((ne11 + r1ptg - 1)/r1ptg), ne12*ne13, 32, nsg, 1);
-    } else if (ggml_metal_op_mul_mat_use_mm(op, props_dev->has_simdgroup_mm)) {
+    } else if (
+        !ggml_is_transposed(op->src[0]) &&
+        !ggml_is_transposed(op->src[1]) &&
+        // for now the matrix-matrix multiplication kernel only works on A14+/M1+ SoCs
+        // AMD GPU and older A-chips will reuse matrix-vector multiplication kernel
+        props_dev->has_simdgroup_mm && ne00 >= 64 && ne11 > ne11_mm_min) {
         //GGML_LOG_INFO("matrix: ne00 = %6d, ne01 = %6d, ne02 = %6d, ne11 = %6d, ne12 = %6d\n", ne00, ne01, ne02, ne11, ne12);
 
         // some Metal matrix data types require aligned pointers
@@ -2613,7 +2622,13 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
     const uint32_t r2 = 1;
     const uint32_t r3 = 1;
 
-    if (ggml_metal_op_mul_mat_id_use_mm(op, props_dev->has_simdgroup_mm)) {
+    // find the break-even point where the matrix-matrix kernel becomes more efficient compared
+    // to the matrix-vector kernel
+    // ne20 = n_used_experts
+    // ne21 = n_rows (batch size)
+    const int ne21_mm_id_min = 32;
+
+    if (props_dev->has_simdgroup_mm && ne00 >= 64 && (ne21 >= ne21_mm_id_min)) {
         // some Metal matrix data types require aligned pointers
         // ref: https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf (Table 2.5)
         //switch (op->src[0]->type) {
@@ -4630,7 +4645,6 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
     const int32_t OW = op->ne[0];
     const int32_t OH = op->ne[1];
     const int32_t OC = op->ne[2];
-    const int32_t N  = op->src[1]->ne[3];
 
     ggml_metal_kargs_conv_transpose_2d args = {
         /*.IC  =*/ IC,
@@ -4643,7 +4657,6 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
         /*.nb0 =*/ nb0,
         /*.nb1 =*/ nb1,
         /*.nb2 =*/ nb2,
-        /*.nb3 =*/ nb3,
     };
 
     auto pipeline = ggml_metal_library_get_pipeline_conv_transpose_2d(lib, op);
@@ -4658,7 +4671,7 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
     const size_t smem = GGML_PAD(KW * KH * sizeof(float), 16);
     ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
 
-    ggml_metal_encoder_dispatch_threadgroups(enc, OW, OH, OC * N, KW, KH, 1);
+    ggml_metal_encoder_dispatch_threadgroups(enc, OW, OH, OC, KW, KH, 1);
 
     return 1;
 }

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

@@ -279,276 +279,6 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = {
     { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
     { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
 
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 128, 128, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 128, 128, 3, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 128, 128, 3, 4 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 192, 128, 1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 192, 128, 2, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 192, 128, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 320, 256, 2, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 320, 256, 3, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 320, 256, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 32, 32, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 32, 32, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 32, 32, 3, 2 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 32, 32, 3, 3 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 64, 64, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 64, 64, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 64, 64, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 1, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 2, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 2, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 3, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 3, 4 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 320, 256, 1, 3 }, { 4, 2 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 576, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, 1, 4 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, 3, 3 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 64, 64, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 96, 96, 2, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 96, 96, 2, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 96, 96, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 96, 96, 3, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, 1, 4 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 192, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 256, 256, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 256, 256, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 320, 256, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 320, 256, 1, 3 }, { 4, 2 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 512, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 512, 512, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 576, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 576, 512, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 32, 32, 1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 32, 32, 2, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 32, 32, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, 1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, 2, 4 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 192, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 192, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 192, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 192, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 128, 1, 3 }, { 4, 2 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 128, 2, 3 }, { 4, 2 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 128, 3, 3 }, { 4, 2 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 256, 256, 1, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 320, 256, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 320, 256, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 320, 256, 1, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 576, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 576, 512, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 32, 32, 1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 32, 32, 2, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 32, 32, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 96, 96, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, 2, 3 }, { 4, 2 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 192, 1, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 128, 1, 3 }, { 4, 2 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 128, 2, 3 }, { 4, 2 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 128, 3, 3 }, { 4, 2 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 256, 256, 1, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 256, 256, 2, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 1, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 2, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 2, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 2, 4 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 3, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 576, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 576, 512, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, 3, 3 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, 1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 96, 96, 2, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 96, 96, 2, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 96, 96, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 96, 96, 3, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 96, 96, 3, 4 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 1, 3 }, { 4, 2 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 2, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 2, 3 }, { 4, 2 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 3, 3 }, { 4, 2 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
-
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 32, 32, 1, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 128, 128, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 128, 128, 1, 1 }, { 1, 1 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 128, 128, 1, 3 }, { 1, 1 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 192, 128, 1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 192, 128, 1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 192, 128, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 192, 128, 1, 3 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 192, 128, 1, 4 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 320, 256, 2, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 320, 256, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 320, 256, 1, 1 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 320, 256, 1, 3 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 320, 256, 1, 4 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 32, 32, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 32, 32, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 32, 32, 2, 3 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 32, 32, 3, 2 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 32, 32, 3, 3 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 64, 64, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 96, 96, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 96, 96, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 96, 96, 1, 4 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 96, 96, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 96, 96, 2, 4 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 96, 96, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 512, 512, 2, 3 }, { 1, 1 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 512, 512, 3, 3 }, { 2, 2 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 576, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 32, 32, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 32, 32, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 32, 32, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 32, 32, 3, 3 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 64, 64, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 64, 64, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 96, 96, 1, 3 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 96, 96, 2, 3 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 96, 96, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 96, 96, 3, 3 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 320, 256, 1, 2 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 320, 256, 1, 4 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 320, 256, 2, 2 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 320, 256, 3, 2 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
-
     { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } },
     { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } },
     { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_F16, 128, 128, 2, 1 }, { 1, 4 } },
@@ -1063,156 +793,6 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = {
     { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 320, 256, 3, 0 }, { 2, 2 } },
     { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } },
     { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 512, 512, 3, 0 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, 1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, 1, 4 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, 2, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, 2, 4 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 64, 64, 3, 2 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 64, 64, 3, 3 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 96, 96, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 96, 96, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 96, 96, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 128, 128, 3, 3 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, 2, 4 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 2, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 3, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 1, 1 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 1, 2 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 1, 3 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 1, 4 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, -1, 1 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, 1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, 1, 4 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, 2, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, 2, 4 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 64, 64, 3, 2 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 96, 96, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 96, 96, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 96, 96, 1, 4 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 96, 96, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 96, 96, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 128, 128, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 128, 128, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 128, 128, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 128, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 128, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 128, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 256, 256, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 256, 256, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 512, 512, 3, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 576, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 576, 512, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 32, 32, 1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 32, 32, 2, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 32, 32, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 96, 96, 1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 96, 96, 2, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 96, 96, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 128, 128, 2, 2 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 128, 128, 2, 3 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 192, 192, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 2, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 256, 256, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 256, 256, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 256, 256, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, -1, 1 }, { 2, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, 1, 2 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, 2, 2 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, 2, 4 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, 3, 2 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, 3, 4 }, { 1, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 3 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 4 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 32, 32, 1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 32, 32, 2, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 32, 32, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, 1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, 2, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, 2, 4 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, 3, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, 3, 4 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, 2, 2 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, 2, 3 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, 3, 2 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, 3, 3 }, { 4, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 192, 192, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 2, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 2, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 256, 256, 1, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 256, 256, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 256, 256, 3, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 320, 256, 1, 1 }, { 2, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 320, 256, 1, 3 }, { 2, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 320, 256, 1, 4 }, { 2, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 320, 256, 3, 1 }, { 2, 2 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 0 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 1 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 2 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 3 }, { 1, 4 } },
-    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 4 }, { 1, 4 } },
     { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } },
     { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } },
     { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 4 }, { 2, 4 } },

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

@@ -558,9 +558,7 @@ static void ggml_backend_metal_event_wait(ggml_backend_t backend, ggml_backend_e
     ggml_metal_event_wait(ctx, ev);
 }
 
-static void ggml_backend_metal_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph, ggml_backend_graph_optimize_params * params) {
-    GGML_UNUSED(params);
-
+static void ggml_backend_metal_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph) {
     ggml_metal_t ctx = (ggml_metal_t)backend->context;
 
     ggml_metal_graph_optimize(ctx, cgraph);

ggml/src/ggml-metal/kernels/conv.metal

@@ -366,8 +366,7 @@ kernel void kernel_conv_transpose_2d(
 
     const int64_t out_x = tgpig[0];
     const int64_t out_y = tgpig[1];
-    const int64_t batch = tgpig[2] / args.OC;
-    const int64_t out_c = tgpig[2] % args.OC;
+    const int64_t out_c = tgpig[2];
 
     const int64_t kw = tpitg[0];
     const int64_t kh = tpitg[1];
@@ -391,7 +390,7 @@ kernel void kernel_conv_transpose_2d(
 
         if (in_x >= args.IW) continue;
 
-        const int64_t input_idx = (args.IW * args.IH) * (args.IC * batch + in_c) + (args.IW) * in_y + in_x;
+        const int64_t input_idx = (args.IW * args.IH) * in_c + (args.IW) * in_y + in_x;
         const int64_t kernel_idx = (args.KH * args.KW * args.OC) * in_c + (args.KH * args.KW) * out_c + (args.KW) * kh + kw;
 
         v += (float)src0[kernel_idx] * src1[input_idx];
@@ -409,7 +408,7 @@ kernel void kernel_conv_transpose_2d(
             total += shared_sum[i];
         }
 
-        device float * dst_ptr = (device float *) (dst + batch*args.nb3 + out_c*args.nb2 + out_y * args.nb1 + out_x*args.nb0);
+        device float * dst_ptr = (device float *) (dst + out_x*args.nb0 + out_y * args.nb1 + out_c*args.nb2);
         dst_ptr[0] = total;
     }
 }

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

@@ -435,7 +435,7 @@ kernel void kernel_mul_mm_id(
         device       char * dst,
         threadgroup  char * shmem [[threadgroup(0)]],
         uint3  tgpig[[threadgroup_position_in_grid]],
-        ushort tiitg[[thread_index_in_threadgroup]],
+        uint   tiitg[[thread_index_in_threadgroup]],
         ushort tiisg[[thread_index_in_simdgroup]],
         ushort sgitg[[simdgroup_index_in_threadgroup]]) {
     threadgroup S0 * sa = (threadgroup S0 *)(shmem);

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

@@ -6059,13 +6059,9 @@ static ggml_backend_opencl_context * ggml_cl_init(ggml_backend_dev_t dev) {
     }
 
 #ifdef GGML_OPENCL_USE_ADRENO_KERNELS
-    // Adreno xmem F16xF32 GEMM, default on adreno, opt out with GGML_OPENCL_ADRENO_XMEM_GEMM=0.
-    // This helps models with f16 attention weights, e.g., gpt-oss-20b-f16
-    {
-        const char * xmem_env = getenv("GGML_OPENCL_ADRENO_XMEM_GEMM");
-        backend_ctx->adreno_xmem_gemm_enabled = backend_ctx->gpu_family == GPU_FAMILY::ADRENO &&
-                                                (xmem_env ? atoi(xmem_env) != 0 : true);
-    }
+    // determine whether to use Adreno xmem GEMM
+    backend_ctx->adreno_xmem_gemm_enabled = getenv("GGML_OPENCL_ADRENO_XMEM_GEMM") != nullptr &&
+                                             backend_ctx->gpu_family == GPU_FAMILY::ADRENO;
 #endif
 
     // determine whether to use large buffer for Adreno
@@ -19538,18 +19534,9 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co
 
     // GEMM using local memory
     // Current BK = 16, so ne00 % 16 == 0
-    //
-    // Certain A7X compiler (E031.41) executes kernel_mul_mm_f32_f32_l4_lm poorly;
-    // matrices with ne11 <= 8 appears OK.
-    // Fallback to the MV style kernels for A7x and ne11 > 8.
-    // Override with GGML_OPENCL_A7X_F32_LM_BYPASS=0.
-    static const char * a7x_f32lm_env    = getenv("GGML_OPENCL_A7X_F32_LM_BYPASS");
-    static const bool   a7x_f32lm_bypass = (a7x_f32lm_env == nullptr || a7x_f32lm_env[0] != '0');
     if (src1t == GGML_TYPE_F32 &&
         ne00 % 16 == 0 &&
-        ne11 > 1 &&
-        !(a7x_f32lm_bypass && src0t == GGML_TYPE_F32 && ne11 > 8 &&
-          backend_ctx->adreno_gen == ADRENO_GPU_GEN::A7X)) {
+        ne11 > 1) {
         switch(src0t) {
             case GGML_TYPE_F32: {
                 kernel = backend_ctx->kernel_mul_mm_f32_f32_l4_lm;

ggml/src/ggml-rpc/CMakeLists.txt

@@ -34,14 +34,10 @@ if (GGML_RPC_RDMA)
         find_library(RDMA_LIB ${RDMA_LIB_NAME} REQUIRED)
     endif()
     target_compile_definitions(ggml-rpc PRIVATE GGML_RPC_RDMA)
+    target_link_libraries(ggml-rpc PRIVATE ${RDMA_LIB})
     if (APPLE)
-        # librdma.dylib only exists on macOS 26.2 and later. Link it weakly so a build made
-        # where it exists still loads where it does not; checked at runtime before use.
-        target_link_options(ggml-rpc PRIVATE "LINKER:-weak_library,${RDMA_LIB}")
         target_compile_definitions(ggml-rpc PRIVATE GGML_RPC_RDMA_APPLE)
         target_sources(ggml-rpc PRIVATE transport-apple.cpp)
-    else()
-        target_link_libraries(ggml-rpc PRIVATE ${RDMA_LIB})
     endif()
     message(STATUS "  RDMA transport enabled (${RDMA_DESC})")
 else()

ggml/src/ggml-rpc/transport-apple.cpp

@@ -8,7 +8,6 @@
 #include <cstdlib>
 #include <cstring>
 #include <string>
-#include <dlfcn.h>
 #include <poll.h>
 #include <sys/socket.h>
 #include <unistd.h>
@@ -185,28 +184,11 @@ static uint8_t rdma_first_active_port(struct ibv_context * ctx, struct ibv_port_
     return 0;
 }
 
-// librdma.dylib is weak-linked, so its symbols are null when it is absent. Nothing may
-// call one before this has returned true.
-static bool rdma_library_present() {
-    static const bool present = [] {
-        void * handle = dlopen("/usr/lib/librdma.dylib", RTLD_LAZY);
-        if (handle == nullptr) {
-            return false;
-        }
-        dlclose(handle);
-        return true;
-    }();
-    return present;
-}
-
 // Called before the endpoints are exchanged: pick the local device facing this
 // peer, create a UC QP and register the frame rings. RDMA is point-to-point, so
 // the device is the one whose GID equals the bootstrap connection's local
 // address, i.e. the one cabled to the peer.
 std::unique_ptr<apple_rdma> apple_rdma::probe(int fd, const uint8_t * target_gid, uint8_t * caps) {
-    if (!rdma_library_present()) {
-        return nullptr;
-    }
     int ndev = 0;
     ibv_device ** devs = ibv_get_device_list(&ndev);
     if (!devs) return nullptr;

ggml/src/ggml-sycl/fattn-common.hpp

@@ -6,7 +6,6 @@
 #include "convert.hpp"
 #include "vecdotq.hpp"
 #include "fattn-buffers.hpp"
-#include "fattn.hpp"
 
 #include "ggml.h"
 
@@ -927,7 +926,6 @@ void launch_fattn(
 
     ggml_sycl_fattn_alloc        K_f16(fbuf.K);
     ggml_sycl_fattn_alloc        V_f16(fbuf.V);
-    const ggml_sycl_fattn_extra  extra = ggml_sycl_fattn_get_extra(dst);
     ggml_sycl_pool_alloc<int>    KV_max(pool);
     ggml_sycl_pool_alloc<float>  dst_tmp(pool);
     ggml_sycl_pool_alloc<sycl::float2> dst_tmp_meta(pool);
@@ -946,11 +944,10 @@ void launch_fattn(
         const size_t bs = ggml_blck_size(K->type);
         const size_t ts = ggml_type_size(K->type);
 
-        sycl::half * K_f16_ptr = extra.K_buffer_ptr ? (sycl::half *) extra.K_buffer_ptr
-                                                    : K_f16.alloc(ggml_nelements(K));
+        K_f16.alloc(ggml_nelements(K));
         if (ggml_is_contiguously_allocated(K)) {
             to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(K->type, dst);
-            to_fp16(K_data, K_f16_ptr, ggml_nelements(K), main_stream);
+            to_fp16(K_data, K_f16.ptr, ggml_nelements(K), main_stream);
 
             nb11 = nb11 * bs * sizeof(sycl::half) / ts;
             nb12 = nb12 * bs * sizeof(sycl::half) / ts;
@@ -961,13 +958,13 @@ void launch_fattn(
             const int64_t s01 = nb11 / ts;
             const int64_t s02 = nb12 / ts;
             const int64_t s03 = nb13 / ts;
-            to_fp16(K_data, K_f16_ptr, K->ne[0], K->ne[1], K->ne[2], K->ne[3], s01, s02, s03, main_stream);
+            to_fp16(K_data, K_f16.ptr, K->ne[0], K->ne[1], K->ne[2], K->ne[3], s01, s02, s03, main_stream);
 
             nb11 = K->ne[0] * sizeof(sycl::half);
             nb12 = K->ne[1] * nb11;
             nb13 = K->ne[2] * nb12;
         }
-        K_data = (char *) K_f16_ptr;
+        K_data = (char *) K_f16.ptr;
     }
 
     if (need_f16_V && V->type != GGML_TYPE_F16) {
@@ -980,12 +977,11 @@ void launch_fattn(
             const size_t bs = ggml_blck_size(V->type);
             const size_t ts = ggml_type_size(V->type);
 
-            sycl::half * V_f16_ptr = extra.V_buffer_ptr ? (sycl::half *) extra.V_buffer_ptr
-                                                        : V_f16.alloc(ggml_nelements(V));
+            V_f16.alloc(ggml_nelements(V));
             if (ggml_is_contiguously_allocated(V)) {
                 to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(V->type, dst);
-                to_fp16(V_data, V_f16_ptr, ggml_nelements(V), main_stream);
-                V_data = (char *) V_f16_ptr;
+                to_fp16(V_data, V_f16.ptr, ggml_nelements(V), main_stream);
+                V_data = (char *) V_f16.ptr;
 
                 nb21 = nb21 * bs * sizeof(sycl::half) / ts;
                 nb22 = nb22 * bs * sizeof(sycl::half) / ts;
@@ -996,13 +992,13 @@ void launch_fattn(
                 const int64_t s01 = nb21 / ts;
                 const int64_t s02 = nb22 / ts;
                 const int64_t s03 = nb23 / ts;
-                to_fp16(V_data, V_f16_ptr, V->ne[0], V->ne[1], V->ne[2], V->ne[3], s01, s02, s03, main_stream);
+                to_fp16(V_data, V_f16.ptr, V->ne[0], V->ne[1], V->ne[2], V->ne[3], s01, s02, s03, main_stream);
 
                 nb21 = V->ne[0] * sizeof(sycl::half);
                 nb22 = V->ne[1] * nb21;
                 nb23 = V->ne[2] * nb22;
             }
-            V_data = (char *) V_f16_ptr;
+            V_data = (char *) V_f16.ptr;
         }
     }
 

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

@@ -14,21 +14,9 @@
 // set minimum query length to treat as prefill (32)
 #define GGML_SYCL_FA_ONEDNN_MIN_Q 32
 
-bool ggml_sycl_fattn_onednn_binds_kv(const ggml_tensor * K, const ggml_tensor * V) {
-    if (K->type != GGML_TYPE_F16 || V->type != GGML_TYPE_F16) {
-        return false;
-    }
-    auto bindable = [](const ggml_tensor * t) {
-        return t->nb[0] == sizeof(sycl::half) && t->nb[1] % sizeof(sycl::half) == 0 &&
-               t->nb[2] % sizeof(sycl::half) == 0 && t->nb[3] % sizeof(sycl::half) == 0;
-    };
-    return bindable(K) && bindable(V);
-}
-
-bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst, bool use_shape_limit) {
+bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) {
 #if !GGML_SYCL_DNNL
     GGML_UNUSED(dst);
-    GGML_UNUSED(use_shape_limit);
     return false;
 #else
     if (!g_ggml_sycl_fa_onednn) {
@@ -56,7 +44,7 @@ bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst, bool use
         if (!k_ok || !v_ok) {
             return false;
         }
-        if (use_shape_limit && (Q->ne[1] < 32 || K->ne[1] < 1024)) {
+        if (Q->ne[1] < 32 || K->ne[1] < 1024) {
             return false;
         }
         for (const ggml_tensor * t : {K, V}) {
@@ -106,7 +94,7 @@ bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst, bool use
         return false;
     }
     // Prefill only.
-    if (use_shape_limit && Q->ne[1] < GGML_SYCL_FA_ONEDNN_MIN_Q) {
+    if (Q->ne[1] < GGML_SYCL_FA_ONEDNN_MIN_Q) {
         return false;
     }
     return true;
@@ -252,16 +240,9 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
     dnnl::engine    eng    = ctx.engine_dnnl(stream);
     dnnl::stream    strm   = ctx.stream_dnnl(stream);
 
-    const ggml_sycl_fattn_extra extra = ggml_sycl_fattn_get_extra(dst);
-
     // Q: always f32 -- copy to dense f16.
-    std::optional<ggml_sycl_pool_alloc<sycl::half>> Qf_pool;
-    sycl::half * Qf_ptr = (sycl::half *) extra.Q_buffer_ptr;
-    if (!Qf_ptr) {
-        Qf_pool.emplace(ctx.pool(), (size_t) H * q * d);
-        Qf_ptr = Qf_pool->get();
-    }
-    cont_to_f16_sycl<float>((const char *) Q->data, Qf_ptr, d, q, H, mb, Q->nb[1], Q->nb[2], Q->nb[3], stream);
+    ggml_sycl_pool_alloc<sycl::half> Qf(ctx.pool(), (size_t) H * q * d);
+    cont_to_f16_sycl<float>((const char *) Q->data, Qf.get(), d, q, H, mb, Q->nb[1], Q->nb[2], Q->nb[3], stream);
 
     // K/V: bind the f16 cache in place. llama.cpp permutes it to [token][head][dim], so its head
     // plane is strided rather than dense, which is what an explicit stride vector expresses.
@@ -272,12 +253,11 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
     std::array<int64_t, 5> v_str = k_str;
     std::optional<ggml_sycl_pool_alloc<sycl::half>> Kf_pool;
     std::optional<ggml_sycl_pool_alloc<sycl::half>> Vf_pool;
-    // Helper: hand out reserved space, or fall back to the pool.
-    auto stage_k = [&](size_t n) { if (extra.K_buffer_ptr) { return (sycl::half *) extra.K_buffer_ptr; }
-                                  Kf_pool.emplace(ctx.pool(), n); return Kf_pool->get(); };
-    auto stage_v = [&](size_t n) { if (extra.V_buffer_ptr) { return (sycl::half *) extra.V_buffer_ptr; }
-                                  Vf_pool.emplace(ctx.pool(), n); return Vf_pool->get(); };
 
+    auto bindable = [](const ggml_tensor * t) {
+        return t->nb[0] == sizeof(sycl::half) && t->nb[1] % sizeof(sycl::half) == 0 &&
+               t->nb[2] % sizeof(sycl::half) == 0 && t->nb[3] % sizeof(sycl::half) == 0;
+    };
     auto elem_strides = [](const ggml_tensor * t) {
         const int64_t s1 = (int64_t) (t->nb[1] / t->nb[0]);
         const int64_t s2 = (int64_t) (t->nb[2] / t->nb[0]);
@@ -286,19 +266,22 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
         return std::array<int64_t, 5>{ s3, s2, s2, s1, 1 };
     };
 
-    if (ggml_sycl_fattn_onednn_binds_kv(K, V)) {
+    if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 && bindable(K) && bindable(V)) {
         K_ptr = (sycl::half *) K->data;
         V_ptr = (sycl::half *) V->data;
         k_str = elem_strides(K);
         v_str = elem_strides(V);
     } else if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16) {
-        K_ptr = stage_k((size_t) Hkv * seq * d);
-        V_ptr = stage_v((size_t) Hkv * seq * d);
-        cont_to_f16_sycl<sycl::half>((const char *) K->data, K_ptr, d, seq, Hkv, mb, K->nb[1], K->nb[2], K->nb[3], stream);
-        cont_to_f16_sycl<sycl::half>((const char *) V->data, V_ptr, d, seq, Hkv, mb, V->nb[1], V->nb[2], V->nb[3], stream);
+        Kf_pool.emplace(ctx.pool(), (size_t) Hkv * seq * d);
+        Vf_pool.emplace(ctx.pool(), (size_t) Hkv * seq * d);
+        cont_to_f16_sycl<sycl::half>((const char *) K->data, Kf_pool->get(), d, seq, Hkv, mb, K->nb[1], K->nb[2], K->nb[3], stream);
+        cont_to_f16_sycl<sycl::half>((const char *) V->data, Vf_pool->get(), d, seq, Hkv, mb, V->nb[1], V->nb[2], V->nb[3], stream);
+        K_ptr = Kf_pool->get();
+        V_ptr = Vf_pool->get();
     } else if (ggml_is_quantized(K->type)) {
         // Quantized K/V: dequant to dense F16 using pool, same lifetime as F16 path.
-        K_ptr = stage_k((size_t) ggml_nelements(K));
+        Kf_pool.emplace(ctx.pool(), ggml_nelements(K));
+        K_ptr = Kf_pool->get();
         {
             const char * K_data = (const char *)K->data;
             const bool k_non_dense = ((int64_t)K->ne[1] * K->nb[1] != K->nb[2]) && K->ne[2] > 1;
@@ -332,7 +315,8 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
         // data pointer), their logical values differ because the quantized
         // elements at different positions/offsets represent different K/V
         // data. Master's F16 path also never aliases K and V.
-        V_ptr = stage_v((size_t) ggml_nelements(V));
+        Vf_pool.emplace(ctx.pool(), ggml_nelements(V));
+        V_ptr = Vf_pool->get();
         {
             const char * V_data = (const char *)V->data;
             const bool v_non_dense = ((int64_t)V->ne[1] * V->nb[1] != V->nb[2]) && V->ne[2] > 1;
@@ -363,10 +347,12 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
         }
     } else {
         // F32: strided copy to dense F16 via cont_to_f16_sycl<float>.
-        K_ptr = stage_k((size_t) ggml_nelements(K));
+        Kf_pool.emplace(ctx.pool(), ggml_nelements(K));
+        K_ptr = Kf_pool->get();
         cont_to_f16_sycl<float>((const char *) K->data, K_ptr, K->ne[0], K->ne[1], K->ne[2], K->ne[3],
                                 K->nb[1], K->nb[2], K->nb[3], stream);
-        V_ptr = stage_v((size_t) ggml_nelements(V));
+        Vf_pool.emplace(ctx.pool(), ggml_nelements(V));
+        V_ptr = Vf_pool->get();
         cont_to_f16_sycl<float>((const char *) V->data, V_ptr, V->ne[0], V->ne[1], V->ne[2], V->ne[3],
                                 V->nb[1], V->nb[2], V->nb[3], stream);
     }
@@ -380,21 +366,11 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
     // instead -- the value is captured into the command, so no host memory has to outlive the
     // call, and the enqueue stays async.
     const sycl::half scale_h = (sycl::half) (1.0f / kq_scale);
-    std::optional<ggml_sycl_pool_alloc<sycl::half>> scbuf;
-    sycl::half * scale_dev = (sycl::half *) extra.scale_buffer_ptr;
-    if (!scale_dev) {
-        scbuf.emplace(ctx.pool(), 1);
-        scale_dev = scbuf->get();
-    }
+    ggml_sycl_pool_alloc<sycl::half> scbuf(ctx.pool(), 1);
+    sycl::half * const scale_dev = scbuf.get();
     stream->single_task([=]() { *scale_dev = scale_h; });
 
-    // f16 contiguous SDPA out [mb,H,q,d]
-    std::optional<ggml_sycl_pool_alloc<sycl::half>> outf_pool;
-    sycl::half * outf_ptr = (sycl::half *) extra.out_buffer_ptr;
-    if (!outf_ptr) {
-        outf_pool.emplace(ctx.pool(), (size_t) H * q * d);
-        outf_ptr = outf_pool->get();
-    }
+    ggml_sycl_pool_alloc<sycl::half> outf(ctx.pool(), (size_t) H * q * d);   // f16 contiguous SDPA out [mb,H,q,d]
 
     // compile once per (device, shape, KV strides), reuse across layers/calls. Stride 2 always
     // repeats stride 1 and stride 4 is always 1, so the key covers every entry that can differ.
@@ -416,7 +392,7 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
     }
 
     auto id2ptr = [&](size_t r) -> void * {
-        if (r == E.id_q)     return Qf_ptr;
+        if (r == E.id_q)     return Qf.get();
         if (r == E.id_k)     return K_ptr;
         if (r == E.id_v)     return V_ptr;
         if (r == E.id_scale) return scale_dev;
@@ -428,10 +404,10 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
     for (auto & lt : E.ins) {
         ti.emplace_back(lt, eng, id2ptr(lt.get_id()));
     }
-    tensor to(E.out, eng, outf_ptr);
+    tensor to(E.out, eng, outf.get());
     E.cp.execute(strm, ti, {to});
 
-    permute_sdpa_out_sycl(outf_ptr, (float *) dst->data, mb, H, q, d, stream);
+    permute_sdpa_out_sycl(outf.get(), (float *) dst->data, mb, H, q, d, stream);
     // Single device needs no sync: the dnnl stream wraps this same in-order queue, so the SDPA
     // serializes with the staging kernels before it and the permute/pool reuse after it. The
     // garbage output formerly blamed on the missing sync here was the scale use-after-return

ggml/src/ggml-sycl/fattn-onednn.hpp

@@ -5,11 +5,7 @@
 
 // Static-only check: fused-XMX oneDNN Graph SDPA path==flash-attn op
 // (f16 KV, no softcap/ALiBi, single stream, tuned head_dim, prefill-sized q.)
-bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst, bool use_shape_limit = true);
-
-// True when the oneDNN path binds an F16 KV cache in place instead of staging a dense copy of
-// it. Depends only on the types and strides of K and V, so the answer holds for every call.
-bool ggml_sycl_fattn_onednn_binds_kv(const ggml_tensor * K, const ggml_tensor * V);
+bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst);
 
 // Run flash attention through oneDNN's fused xmx SDPA
 // execute the cached SDPA partition, write the f32 dst. Falls back to the TILE kernel on any failure.

ggml/src/ggml-sycl/fattn.cpp

@@ -378,76 +378,3 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst
 bool ggml_sycl_flash_attn_ext_supported(int device, const ggml_tensor * dst) {
     return ggml_sycl_get_best_fattn_kernel(device, dst) != BEST_FATTN_KERNEL_NONE;
 }
-
-static uintptr_t ggml_sycl_fattn_reserve_halves(ggml_sycl_fattn_extra & extra, size_t n_halves) {
-    if (n_halves == 0) {
-        return 0;
-    }
-    extra.end = GGML_PAD(extra.end, SYCL_BUFFER_ALIGNMENT);
-    const uintptr_t block = extra.end;
-    extra.end += n_halves * sizeof(sycl::half);
-    return block;
-}
-
-ggml_sycl_fattn_extra ggml_sycl_fattn_get_extra(const ggml_tensor * dst) {
-    ggml_sycl_fattn_extra extra;
-
-    extra.end = (uintptr_t) dst->data + ggml_nbytes(dst);
-
-    if (dst->op != GGML_OP_FLASH_ATTN_EXT) {
-        return extra;
-    }
-
-    const ggml_tensor * Q = dst->src[0];
-    const ggml_tensor * K = dst->src[1];
-    const ggml_tensor * V = dst->src[2];
-    if (!Q || !K || !V) {
-        return extra;
-    }
-
-    const int64_t d = K->ne[0];
-    const int64_t H = Q->ne[2];
-    const int64_t q = Q->ne[1];
-
-    // calculate the worst-case memory consumption across all kernels
-    const bool onednn_supported = ggml_sycl_flash_attn_ext_onednn_supported(dst, /* use_shape_limit */ false);
-
-    const bool tile_needs_K = K->type != GGML_TYPE_F16;
-    const bool tile_needs_V = V->type != GGML_TYPE_F16;
-
-    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));
-
-    size_t need_K = 0, need_V = 0, need_Q = 0, need_out = 0, need_scale = 0;
-    if (onednn_supported) {
-        need_Q     = (size_t) H * q * d;
-        need_out   = (size_t) H * q * d;
-        need_scale = 1;
-        // an f16 cache is bound in place, so it needs no staging copy
-        if (!ggml_sycl_fattn_onednn_binds_kv(K, V)) {
-            need_K = (size_t) ggml_nelements(K);
-            need_V = (size_t) ggml_nelements(V);
-        }
-    }
-    if (tile_needs_K) {
-        need_K = std::max(need_K, (size_t) ggml_nelements(K));
-    }
-    if (tile_needs_V) {
-        need_V = std::max(need_V, (size_t) ggml_nelements(V));
-    }
-
-    extra.Q_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_Q);
-    extra.K_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_K);
-    extra.V_buffer_ptr = (V_is_K_view && !onednn_supported && need_V)
-                       ? extra.K_buffer_ptr
-                       : ggml_sycl_fattn_reserve_halves(extra, need_V);
-    extra.scale_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_scale);
-    extra.out_buffer_ptr   = ggml_sycl_fattn_reserve_halves(extra, need_out);
-
-    return extra;
-}
-
-size_t ggml_sycl_flash_attn_ext_get_alloc_size(const ggml_tensor * dst) {
-    const ggml_sycl_fattn_extra extra = ggml_sycl_fattn_get_extra(dst);
-    return (size_t) (extra.end - (uintptr_t) dst->data);
-}

ggml/src/ggml-sycl/fattn.hpp

@@ -19,24 +19,6 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst
 
 bool ggml_sycl_flash_attn_ext_supported(int device, const ggml_tensor * dst);
 
-// Scratch that flash attention needs beyond the output tensor
-struct ggml_sycl_fattn_extra {
-    uintptr_t K_buffer_ptr     = 0;   // F16 copy of the K cache
-    uintptr_t V_buffer_ptr     = 0;   // F16 copy of the V cache
-    uintptr_t Q_buffer_ptr     = 0;   // dense F16 copy of Q, oneDNN only
-    uintptr_t scale_buffer_ptr = 0;   // the softmax scale as an F16 scalar, oneDNN only
-    uintptr_t out_buffer_ptr   = 0;   // F16 SDPA output before conversion to F32, oneDNN only
-    uintptr_t end              = 0;   // one past the last reserved byte; sizes the allocation
-};
-
-// ggml_sycl_fattn_get_extra() is the single source of truth for the layout: it both sizes
-// the reservation and hands out the pointers, so the two cannot disagree.
-// Each field is the address of one reserved block, or 0 if that block was not reserved,
-// in which case the caller allocates from the scratch pool instead.
-ggml_sycl_fattn_extra ggml_sycl_fattn_get_extra(const ggml_tensor * dst);
-
-size_t ggml_sycl_flash_attn_ext_get_alloc_size(const ggml_tensor * dst);
-
 void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
 
 #endif // GGML_SYCL_FATTN_HPP

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

@@ -955,10 +955,7 @@ static size_t ggml_backend_sycl_buffer_type_get_max_size(ggml_backend_buffer_typ
 }
 
 static size_t ggml_backend_sycl_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) {
-    // Reserve the additional scratch so it's visible to the graph allocator
-    size_t size = tensor->op == GGML_OP_FLASH_ATTN_EXT
-        ? ggml_sycl_flash_attn_ext_get_alloc_size(tensor)
-        : ggml_nbytes(tensor);
+    size_t size = ggml_nbytes(tensor);
     int64_t ne0 = tensor->ne[0];
 
     if (ggml_is_quantized(tensor->type)) {
@@ -2402,138 +2399,7 @@ static void argsort_f32_i32_sycl(const float *x, int *dst, const int ncols,
     }
 }
 
-// Scan and block merge, shared by every launch shape below so a partitioned row uses the
-// same insertion order as an unpartitioned one.
-//
-// src_map  != nullptr: report src_map[col] instead of col, so a merge pass can carry the
-//                      original column index through.
-// out_vals != nullptr: also emit the k winning values, for a later merge pass.
-// swap01:              emit in the output order the single-pass path uses.
-static void top_k_scan_merge_f32(
-    const float *   src_vals,
-    const int32_t * src_map,
-    const int       begin,
-    const int       end,
-    const int       k,
-    const int       block_size,
-    float *         shared_vals,
-    int *           shared_idx,
-    float *         out_vals,
-    int32_t *       out_idx,
-    const bool      swap01,
-    const sycl::nd_item<1> & item_ct1
-) {
-    const int tid = item_ct1.get_local_id(0);
-
-    // The running top-k lives in SLM (shared local memory) rather than a private array:
-    // an array indexed by a runtime position cannot be register-allocated, so a private
-    // one lands in scratch, i.e. device memory, and insertion is this kernel's dominant
-    // cost.
-    //
-    // Lane-strided (lv[i * block_size]) rather than lane-blocked (lv[i]) so a given i is
-    // contiguous across lanes; a k-strided layout would put every lane of a shift step in
-    // the same SLM bank.
-    float * lv = shared_vals + tid;
-    int * li = shared_idx + tid;
-
-    for (int i = 0; i < k; i++) {
-        lv[i * block_size] = -FLT_MAX;
-        li[i * block_size] = -1;
-    }
-
-    // The k-th best, cached in a register. The reject test is taken for the large
-    // majority of elements scanned, and in that case touches no memory.
-    float kth = -FLT_MAX;
-
-    for (int col = begin + tid; col < end; col += block_size) {
-        float val = src_vals[col];
-
-        if (val > kth) {
-            int pos = k - 1;
-            while (pos > 0 && val > lv[(pos - 1) * block_size]) {
-                pos--;
-            }
-
-            for (int i = k - 1; i > pos; i--) {
-                lv[i * block_size] = lv[(i - 1) * block_size];
-                li[i * block_size] = li[(i - 1) * block_size];
-            }
-            lv[pos * block_size] = val;
-            li[pos * block_size] = src_map ? src_map[col] : col;
-
-            kth = lv[(k - 1) * block_size];
-        }
-    }
-
-    item_ct1.barrier(sycl::access::fence_space::local_space);
-
-    if (tid != 0) {
-        return;
-    }
-
-    // Same treatment for the merge accumulator, past the per-lane region.
-    float * fv = shared_vals + (size_t) k * block_size;
-    int * fi = shared_idx + (size_t) k * block_size;
-
-    for (int i = 0; i < k; i++) {
-        fv[i] = -FLT_MAX;
-        fi[i] = -1;
-    }
-
-    float fkth = -FLT_MAX;
-
-    // Candidates are visited in the same (t, i) order as before, so tie-breaking is
-    // unchanged.
-    for (int t = 0; t < block_size; t++) {
-        for (int i = 0; i < k; i++) {
-            float val = shared_vals[i * block_size + t];
-
-            if (val <= fkth) {
-                // Lane t's list is sorted descending, so once one of its entries loses
-                // to the k-th best, every later entry loses too. fkth only rises, so
-                // that stays true for the rest of the merge. This turns the merge from
-                // block_size*k steps into roughly block_size plus the candidates
-                // accepted.
-                break;
-            }
-
-            int idx = shared_idx[i * block_size + t];
-
-            int pos = k - 1;
-            while (pos > 0 && val > fv[pos - 1]) {
-                pos--;
-            }
-
-            for (int j = k - 1; j > pos; j--) {
-                fv[j] = fv[j - 1];
-                fi[j] = fi[j - 1];
-            }
-            fv[pos] = val;
-            fi[pos] = idx;
-
-            fkth = fv[k - 1];
-        }
-    }
-
-    if (out_vals) {
-        for (int i = 0; i < k; i++) {
-            out_vals[i] = fv[i];
-        }
-    }
-
-    for (int i = 0; i < k; i++) {
-        out_idx[i] = fi[i];
-    }
-
-    if (swap01 && k > 1) {
-        int32_t temp = out_idx[0];
-        out_idx[0] = out_idx[1];
-        out_idx[1] = temp;
-    }
-}
-
 static void top_k_f32_sycl(
-    ggml_backend_sycl_context & ctx,
     const float * src,
     int32_t * dst_indices,
     const int64_t ncols,
@@ -2541,107 +2407,98 @@ static void top_k_f32_sycl(
     const int k,
     dpct::queue_ptr main_stream
 ) {
-    // A row is scanned by exactly one work-group, so a vocabulary-sized row leaves the
-    // rest of the device idle. What the scan is short of is memory requests in flight,
-    // not bandwidth or per-request latency, so lanes in flight is the lever: split the
-    // row across independent work-groups, have each emit its partition's top-k, and
-    // merge those nsplit*k candidates in a second launch.
-    //
-    // split_block trades parallelism against SLM residency. Its cost is
-    // (split_block + 1) * k * 8 bytes of SLM per group, so at the k <= 32 ceiling 128
-    // lanes need about 33 KB, which leaves a single resident group per Xe-core. Revisit
-    // if the supported k ever grows.
-    constexpr int split_block = 128;
-    constexpr int max_splits  = 128;
-    constexpr int min_cols    = 8192;
-
-    int nsplit = 1;
-    if (ncols >= min_cols) {
-        // A partition is then always >= split_block = 128 columns, hence always more than
-        // the k <= 32 ceiling, so no pass is ever padded with -FLT_MAX sentinels.
-        const int64_t want = ncols / split_block;
-        nsplit = (int) (want > max_splits ? max_splits : want);
-    }
-
-    if (nsplit > 1) {
-        const int nchunk = (int) ((ncols + nsplit - 1) / nsplit);
-        const size_t ncand = (size_t) nrows * nsplit * k;
-
-        ggml_sycl_pool_alloc<float>   part_vals(ctx.pool(), ncand);
-        ggml_sycl_pool_alloc<int32_t> part_idx(ctx.pool(), ncand);
-
-        float *   pv = part_vals.get();
-        int32_t * pi = part_idx.get();
-
-        const sycl::range<1> block_dims(split_block);
-
-        main_stream->submit([&](sycl::handler &cgh) {
-            sycl::local_accessor<float, 1> shared_vals(sycl::range<1>((split_block + 1) * k), cgh);
-            sycl::local_accessor<int, 1> shared_idx(sycl::range<1>((split_block + 1) * k), cgh);
-
-            cgh.parallel_for(
-                sycl::nd_range<1>(sycl::range<1>(nrows * nsplit) * block_dims, block_dims),
-                [=](sycl::nd_item<1> item_ct1) {
-                    const int grp = item_ct1.get_group(0);
-                    const int row = grp / nsplit;
-                    const int part = grp % nsplit;
-
-                    const int begin = part * nchunk;
-                    int end = begin + nchunk;
-                    if (end > (int) ncols) {
-                        end = (int) ncols;
-                    }
-
-                    top_k_scan_merge_f32(
-                        src + (int64_t) row * ncols, nullptr, begin, end, k, split_block,
-                        shared_vals.get_multi_ptr<sycl::access::decorated::no>().get(),
-                        shared_idx.get_multi_ptr<sycl::access::decorated::no>().get(),
-                        pv + (size_t) grp * k, pi + (size_t) grp * k, false, item_ct1);
-                });
-        });
-
-        main_stream->submit([&](sycl::handler &cgh) {
-            sycl::local_accessor<float, 1> shared_vals(sycl::range<1>((split_block + 1) * k), cgh);
-            sycl::local_accessor<int, 1> shared_idx(sycl::range<1>((split_block + 1) * k), cgh);
-
-            cgh.parallel_for(
-                sycl::nd_range<1>(sycl::range<1>(nrows) * block_dims, block_dims),
-                [=](sycl::nd_item<1> item_ct1) {
-                    const int row = item_ct1.get_group(0);
-                    const size_t off = (size_t) row * nsplit * k;
-
-                    top_k_scan_merge_f32(
-                        pv + off, pi + off, 0, nsplit * k, k, split_block,
-                        shared_vals.get_multi_ptr<sycl::access::decorated::no>().get(),
-                        shared_idx.get_multi_ptr<sycl::access::decorated::no>().get(),
-                        nullptr, dst_indices + (int64_t) row * k, true, item_ct1);
-                });
-        });
-
-        return;
-    }
-
     const int block_size = 128;
 
     const sycl::range<1> block_dims(block_size);
     const sycl::range<1> grid_dims(nrows);
 
     main_stream->submit([&](sycl::handler &cgh) {
-        sycl::local_accessor<float, 1> shared_vals(sycl::range<1>((block_size + 1) * k), cgh);
-        sycl::local_accessor<int, 1> shared_idx(sycl::range<1>((block_size + 1) * k), cgh);
+        sycl::local_accessor<float, 1> shared_vals(sycl::range<1>(block_size * k), cgh);
+        sycl::local_accessor<int, 1> shared_idx(sycl::range<1>(block_size * k), cgh);
 
         cgh.parallel_for(
             sycl::nd_range<1>(grid_dims * block_dims, block_dims),
             [=](sycl::nd_item<1> item_ct1) {
                 const int row = item_ct1.get_group(0);
+                const int tid = item_ct1.get_local_id(0);
 
                 if (row >= nrows) return;
 
-                top_k_scan_merge_f32(
-                    src + (int64_t) row * ncols, nullptr, 0, (int) ncols, k, block_size,
-                    shared_vals.get_multi_ptr<sycl::access::decorated::no>().get(),
-                    shared_idx.get_multi_ptr<sycl::access::decorated::no>().get(),
-                    nullptr, dst_indices + (int64_t) row * k, true, item_ct1);
+                const float * src_row = src + row * ncols;
+                int32_t * dst_idx_row = dst_indices + row * k;
+
+                float local_vals[32];
+                int local_idx[32];
+
+                for (int i = 0; i < k; i++) {
+                    local_vals[i] = -FLT_MAX;
+                    local_idx[i] = -1;
+                }
+
+                for (int col = tid; col < ncols; col += block_size) {
+                    float val = src_row[col];
+
+                    if (val > local_vals[k-1]) {
+                        int pos = k - 1;
+                        while (pos > 0 && val > local_vals[pos - 1]) {
+                            pos--;
+                        }
+
+                        for (int i = k - 1; i > pos; i--) {
+                            local_vals[i] = local_vals[i - 1];
+                            local_idx[i] = local_idx[i - 1];
+                        }
+                        local_vals[pos] = val;
+                        local_idx[pos] = col;
+                    }
+                }
+
+                for (int i = 0; i < k; i++) {
+                    shared_vals[tid * k + i] = local_vals[i];
+                    shared_idx[tid * k + i] = local_idx[i];
+                }
+                item_ct1.barrier(sycl::access::fence_space::local_space);
+
+                if (tid == 0) {
+                    float final_vals[32];
+                    int final_idx[32];
+
+                    for (int i = 0; i < k; i++) {
+                        final_vals[i] = -FLT_MAX;
+                        final_idx[i] = -1;
+                    }
+
+                    for (int t = 0; t < block_size; t++) {
+                        for (int i = 0; i < k; i++) {
+                            float val = shared_vals[t * k + i];
+                            int idx = shared_idx[t * k + i];
+
+                            if (val > final_vals[k-1]) {
+                                int pos = k - 1;
+                                while (pos > 0 && val > final_vals[pos - 1]) {
+                                    pos--;
+                                }
+
+                                for (int j = k - 1; j > pos; j--) {
+                                    final_vals[j] = final_vals[j - 1];
+                                    final_idx[j] = final_idx[j - 1];
+                                }
+                                final_vals[pos] = val;
+                                final_idx[pos] = idx;
+                            }
+                        }
+                    }
+
+                    for (int i = 0; i < k; i++) {
+                        dst_idx_row[i] = final_idx[i];
+                    }
+
+                    if (k > 1) {
+                        int32_t temp = dst_idx_row[0];
+                        dst_idx_row[0] = dst_idx_row[1];
+                        dst_idx_row[1] = temp;
+                    }
+                }
             });
     });
 }
@@ -3042,7 +2899,7 @@ static void ggml_sycl_op_top_k(ggml_backend_sycl_context & ctx, ggml_tensor * ds
     GGML_ASSERT(k > 0 && k <= 32);
     GGML_ASSERT(k <= ncols);
 
-    top_k_f32_sycl(ctx, src0_dd, dst_dd, ncols, nrows, k, main_stream);
+    top_k_f32_sycl(src0_dd, dst_dd, ncols, nrows, k, main_stream);
 }
 
 inline void ggml_sycl_op_argmax(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {

ggml/src/ggml-virtgpu/ggml-backend.cpp

@@ -17,8 +17,7 @@ static ggml_status ggml_backend_remoting_graph_compute(ggml_backend_t backend, g
     return apir_backend_graph_compute(gpu, cgraph);
 }
 
-static void ggml_backend_remoting_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph, ggml_backend_graph_optimize_params * params) {
-    UNUSED(params);
+static void ggml_backend_remoting_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph) {
     virtgpu * gpu = DEV_TO_GPU(backend->device);
 #if true
     UNUSED(gpu);

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

@@ -1347,8 +1347,7 @@ struct vk_mat_mat_id_push_constants {
     uint32_t stride_a; uint32_t stride_b; uint32_t stride_d;
     uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d;
     uint32_t nei0; uint32_t nei1; uint32_t nbi1; uint32_t ne11;
-    uint32_t n_experts;
-    uint32_t hoist_row_ids;
+    uint32_t padded_N;
 };
 struct vk_mat_vec_id_push_constants {
     uint32_t ncols;
@@ -1429,10 +1428,6 @@ struct vk_op_count_experts_push_constants {
     uint32_t nb00;
     uint32_t nb01;
     uint32_t a_offset;
-    uint32_t n_experts;
-    uint32_t hoist_row_ids;
-    uint32_t ne00mp;
-    uint32_t ne00L;
 };
 
 struct vk_op_glu_push_constants {
@@ -1611,10 +1606,6 @@ template <> void init_pushconst_fastdiv(vk_op_glu_push_constants &p) {
     init_fastdiv_values(p.ne20,                p.ne2_0mp,      p.ne2_0L);
 }
 
-template <> void init_pushconst_fastdiv(vk_op_count_experts_push_constants &p) {
-    init_fastdiv_values(p.ne00, p.ne00mp, p.ne00L);
-}
-
 struct vk_op_binary_push_constants {
     uint32_t ne;
     uint32_t ne00; uint32_t ne01; uint32_t ne02; uint32_t ne03; uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03;
@@ -2402,8 +2393,9 @@ struct ggml_backend_vk_context {
     // Cache most recent tensor that was converted into prealloc_y, and what pipeline it used to convert.
     vk_pipeline_struct * prealloc_y_last_pipeline_used {};
     const ggml_tensor * prealloc_y_last_tensor_used {};
-    // True when the K dimension in prealloc_y is padded.
-    bool prealloc_y_last_k_padded {};
+    // True when prealloc_y holds the padded fp16 layout used by the coopmat2 B decode-vector callback.
+    // If false, then it's contiguous.
+    bool prealloc_y_last_decode_vector_staging {};
 
     // Track which nodes have been used since the last sync, and whether they were written to
     std::vector<const ggml_tensor *> unsynced_nodes_written;
@@ -5847,11 +5839,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
 
     ggml_vk_create_pipeline(device, device->pipeline_count_equal_i32, "count_equal_i32", count_equal_i32_len, count_equal_i32_data, "main", 3, sizeof(vk_op_push_constants), {512, 1, 1}, { device->subgroup_size }, 1);
 
-    if (device->subgroup_arithmetic && device->subgroup_require_full_support) {
-        ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_subgroup_len, count_experts_subgroup_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true, true);
-    } else {
-        ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_len, count_experts_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true);
-    }
+    ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_len, count_experts_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true);
 
     for (auto &s : device->pipeline_solve_tri_f32) {
         const vk_solve_tri_pipeline_state &state = s.first;
@@ -8982,13 +8970,13 @@ static void ggml_vk_matmul_id(
         uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d,
         uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d,
         uint32_t n_as, uint32_t nei0, uint32_t nei1, uint32_t nbi1, uint32_t ne11,
-        bool hoist_row_ids) {
+        uint32_t padded_n) {
     VK_LOG_DEBUG("ggml_vk_matmul_id(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), ids: (" << ids.buffer->buffer << ", " << ids.offset << ", " << ids.size << "), expert_count: (" << expert_count_buf.buffer->buffer << ", " << expert_count_buf.offset << ", " << expert_count_buf.size << "), " <<
         "m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", " <<
         "batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", " <<
         "n_as: " << n_as << ", nei0: " << nei0 << ", nei1: " << nei1 << ", nbi1: " << nbi1 << ", ne11: " << ne11 << ")");
     const vk_mat_mat_id_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d,
-                                              nei0, nei1, nbi1, ne11, n_as, uint32_t(hoist_row_ids) };
+                                              nei0, nei1, nbi1, ne11, padded_n };
     ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d, ids, expert_count_buf }, pc, { m, nei1, n_as });
 }
 
@@ -9453,27 +9441,27 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub
     if (y_non_contig) {
         if (ctx->prealloc_y_last_pipeline_used != to_fp16_vk_1.get() ||
             ctx->prealloc_y_last_tensor_used != src1 ||
-            ctx->prealloc_y_last_k_padded) {
+            ctx->prealloc_y_last_decode_vector_staging) {
             if (ctx->prealloc_y_need_sync) {
                 ggml_vk_sync_buffers(ctx, subctx);
             }
             ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, ggml_vk_subbuffer(ctx, d_Qy, qy_buf_offset), ggml_vk_subbuffer(ctx, d_Y, 0));
             ctx->prealloc_y_last_pipeline_used = to_fp16_vk_1.get();
             ctx->prealloc_y_last_tensor_used = src1;
-            ctx->prealloc_y_last_k_padded = false;
+            ctx->prealloc_y_last_decode_vector_staging = false;
         }
     }
     if (quantize_y) {
         if (ctx->prealloc_y_last_pipeline_used != to_q8_1.get() ||
             ctx->prealloc_y_last_tensor_used != src1 ||
-            ctx->prealloc_y_last_k_padded) {
+            ctx->prealloc_y_last_decode_vector_staging) {
             if (ctx->prealloc_y_need_sync) {
                 ggml_vk_sync_buffers(ctx, subctx);
             }
             ggml_vk_quantize_q8_1(ctx, subctx, ggml_vk_subbuffer(ctx, d_Qy, qy_buf_offset), ggml_vk_subbuffer(ctx, d_Y, 0), y_ne);
             ctx->prealloc_y_last_pipeline_used = to_q8_1.get();
             ctx->prealloc_y_last_tensor_used = src1;
-            ctx->prealloc_y_last_k_padded = false;
+            ctx->prealloc_y_last_decode_vector_staging = false;
         }
     }
 
@@ -9732,27 +9720,27 @@ static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context&
         GGML_ASSERT(y_sz == ggml_type_size(src1->type) * y_ne);
         if (ctx->prealloc_y_last_pipeline_used != to_fp16_vk_1.get() ||
             ctx->prealloc_y_last_tensor_used != src1 ||
-            ctx->prealloc_y_last_k_padded) {
+            ctx->prealloc_y_last_decode_vector_staging) {
             if (ctx->prealloc_y_need_sync) {
                 ggml_vk_sync_buffers(ctx, subctx);
             }
             ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, d_Qy, d_Y);
             ctx->prealloc_y_last_pipeline_used = to_fp16_vk_1.get();
             ctx->prealloc_y_last_tensor_used = src1;
-            ctx->prealloc_y_last_k_padded = false;
+            ctx->prealloc_y_last_decode_vector_staging = false;
         }
     }
     if (quantize_y) {
         if (ctx->prealloc_y_last_pipeline_used != to_q8_1.get() ||
             ctx->prealloc_y_last_tensor_used != src1 ||
-            ctx->prealloc_y_last_k_padded) {
+            ctx->prealloc_y_last_decode_vector_staging) {
             if (ctx->prealloc_y_need_sync) {
                 ggml_vk_sync_buffers(ctx, subctx);
             }
             ggml_vk_quantize_q8_1(ctx, subctx, d_Qy, d_Y, y_ne);
             ctx->prealloc_y_last_pipeline_used = to_q8_1.get();
             ctx->prealloc_y_last_tensor_used = src1;
-            ctx->prealloc_y_last_k_padded = false;
+            ctx->prealloc_y_last_decode_vector_staging = false;
         }
     }
 
@@ -10174,12 +10162,6 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
     // const uint64_t ne23 = dst->ne[3];
 
     const uint64_t n_as = ne02;
-    // n_as counts, n_as offsets, one total, then one packed row id per (expert, token).
-    // Hoisting requires 16-bit indices for the packing and a table that fits one binding.
-    const uint64_t hoisted_row_id_words = 2 * n_as + 1 + nei0 * nei1;
-    const bool hoist_row_ids = n_as <= 256 && nei0 <= 0xffff && nei1 <= 0xffff &&
-                                hoisted_row_id_words * sizeof(uint32_t) <=
-                                    ctx->device->properties.limits.maxStorageBufferRange;
 
     ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context;
     ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context;
@@ -10232,6 +10214,8 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
                               (src0->type == GGML_TYPE_BF16 && src1->type != GGML_TYPE_BF16) ||
                               !ggml_vk_dim01_contiguous(src1);
 
+    const uint32_t y_staged_row_stride = y_decode_vector_staging ? (uint32_t)ggml_vk_align_size(ne10, 4) : (uint32_t)ne10;
+
     const bool y_f32_kernel = src1->type == GGML_TYPE_F32 && !y_non_contig;
 
     bool quantize_y = ctx->device->integer_dot_product && src1->type == GGML_TYPE_F32 && ggml_is_contiguous(src1) && !y_non_contig && (ne11 * ne10) % 4 == 0;
@@ -10246,25 +10230,19 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
     }
 
     const bool qx_needs_dequant = mmp == nullptr || x_non_contig;
-    bool qy_needs_dequant = !quantize_y && ((src1->type != f16_type && !y_f32_kernel) || y_non_contig);
+    const bool qy_needs_dequant = !quantize_y && ((src1->type != f16_type && !y_f32_kernel) || y_non_contig);
 
     if (qx_needs_dequant) {
         // Fall back to dequant + f16 mulmat
         mmp = ggml_vk_get_mul_mat_mat_id_pipeline(ctx, f16_type, y_f32_kernel ? GGML_TYPE_F32 : f16_type, (ggml_prec)dst->op_params[0]);
     }
 
+    // Not implemented
+    GGML_ASSERT(y_non_contig || !qy_needs_dequant);  // NOLINT
+
     const ggml_type effective_src1_type = quantize_y ? GGML_TYPE_Q8_1 : (y_f32_kernel ? GGML_TYPE_F32 : src1->type);
 
     const uint32_t kpad = quantize_y ? 0 : ggml_vk_align_size(ne10, ggml_vk_guess_matmul_id_pipeline_align(ctx, mmp, ne01, nei1, qx_needs_dequant ? f16_type : src0->type, effective_src1_type));
-    // Coopmat2 MUL_MAT_ID BK specialization constants in ggml_vk_load_shaders are at most 64.
-    const uint32_t y_staged_row_stride = ctx->device->coopmat2 && !quantize_y ? ggml_vk_align_size(ne10, 64) : ne10;
-    const bool y_needs_k_padding = ne10 != y_staged_row_stride;
-    const bool y_needs_reformat = y_non_contig || y_needs_k_padding;
-    qy_needs_dequant = qy_needs_dequant || y_needs_k_padding;
-
-    // Not implemented
-    GGML_ASSERT(y_needs_reformat || !qy_needs_dequant);  // NOLINT
-
     const bool aligned = !quantize_y && ne10 == kpad && ne01 > 8 && nei1 > 8;
 
     vk_pipeline pipeline = ggml_vk_guess_matmul_id_pipeline(ctx, mmp, ne01, nei1, aligned, qx_needs_dequant ? f16_type : src0->type, effective_src1_type);
@@ -10272,8 +10250,10 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
     if (ggml_nbytes(src0) > ctx->device->properties.limits.maxStorageBufferRange) {
         pipeline = ggml_vk_get_64b_indexing_pipeline(ctx, pipeline);
     }
+    // Reserve extra storage in the N dimension for the Y matrix, so we can avoid bounds-checking
+    uint32_t padded_n = qy_needs_dequant ? ROUNDUP_POW2(ne11, pipeline->wg_denoms[1]) :ne11;
     const uint64_t x_ne = ggml_nelements(src0);
-    const uint64_t y_ne = (uint64_t)y_staged_row_stride * ne11 * ne12 * ne13;
+    const uint64_t y_ne = (uint64_t)y_staged_row_stride * padded_n * ne12 * ne13;
     const uint64_t d_ne = ggml_nelements(dst);
 
     const uint64_t qx_sz = ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type);
@@ -10292,7 +10272,7 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
         y_staged_dst.type = f16_type;
         y_staged_dst.nb[0] = ggml_type_size(f16_type);
         y_staged_dst.nb[1] = y_staged_dst.nb[0] * y_staged_row_stride;
-        y_staged_dst.nb[2] = y_staged_dst.nb[1] * ne11;
+        y_staged_dst.nb[2] = y_staged_dst.nb[1] * padded_n;
         y_staged_dst.nb[3] = y_staged_dst.nb[2] * y_staged_dst.ne[2];
         return y_staged_dst;
     };
@@ -10302,10 +10282,10 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
     } else {
         to_fp16_vk_0 = ggml_vk_get_to_fp16(ctx, src0->type);
     }
-    if (y_needs_reformat) {
+    if (y_non_contig) {
         ggml_tensor y_staged_dst;
         const ggml_tensor * y_staged_dst_ptr = nullptr;
-        if (y_needs_k_padding) {
+        if (y_decode_vector_staging) {
             y_staged_dst = make_y_staged_dst();
             y_staged_dst_ptr = &y_staged_dst;
         }
@@ -10322,8 +10302,7 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
     }
     vk_pipeline count_experts = ctx->device->pipeline_count_experts;
 
-    const size_t expert_data_size = sizeof(uint32_t) *
-        (hoist_row_ids ? hoisted_row_id_words : n_as);
+    uint32_t expert_count_size = sizeof(uint32_t) * n_as;
 
     {
         if (
@@ -10339,8 +10318,8 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
             ctx->prealloc_size_y = y_sz;
             ggml_vk_preallocate_buffers(ctx, subctx);
         }
-        if (ctx->prealloc_size_split_k < expert_data_size) {
-            ctx->prealloc_size_split_k = expert_data_size;
+        if (ctx->prealloc_size_split_k < expert_count_size) {
+            ctx->prealloc_size_split_k = expert_count_size;
             ggml_vk_preallocate_buffers(ctx, subctx);
         }
 
@@ -10406,23 +10385,18 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
         }
     }
     // Count how many times each expert is used
-    vk_subbuffer expert_count_buf = { ctx->prealloc_split_k, 0, expert_data_size };
+    vk_subbuffer expert_count_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0);
     if (ctx->prealloc_split_k_need_sync) {
         ggml_vk_sync_buffers(ctx, subctx);
     }
     {
-        vk_op_count_experts_push_constants pc = { (uint32_t)nei0,
+        const std::vector<uint32_t> pc = { (uint32_t)nei0,
                                            (uint32_t)nei1,
                                            (uint32_t)(nbi0 / ggml_type_size(ids->type)),
                                            (uint32_t)(nbi1 / ggml_type_size(ids->type)),
-                                           (uint32_t)(get_misalign_bytes(ctx, ids) / ggml_type_size(ids->type)),
-                                           (uint32_t)n_as,
-                                           uint32_t(hoist_row_ids),
-                                           0, 0 };
-        init_pushconst_fastdiv(pc);
+                                           (uint32_t)(get_misalign_bytes(ctx, ids) / ggml_type_size(ids->type)) };
         ggml_vk_dispatch_pipeline(ctx, subctx, count_experts,
-            { vk_subbuffer{ d_ids, ids_buf_offset, ids_sz }, expert_count_buf }, pc,
-            { hoist_row_ids ? 1u : (uint32_t)n_as, 1, 1});
+            { vk_subbuffer{ d_ids, ids_buf_offset, ids_sz }, expert_count_buf }, pc, { (uint32_t)n_as, 1, 1});
     }
 
     if (x_non_contig) {
@@ -10432,18 +10406,14 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
         ggml_vk_dispatch_pipeline(ctx, subctx, to_fp16_vk_0,
             { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz }, vk_subbuffer{ d_X, 0, x_sz } }, pc, { (uint32_t)x_ne, 1, 1});
     }
-    if (y_needs_reformat) {
+    if (y_non_contig) {
         if (ctx->prealloc_y_last_pipeline_used != to_fp16_vk_1.get() ||
             ctx->prealloc_y_last_tensor_used != src1 ||
-            ctx->prealloc_y_last_k_padded != y_needs_k_padding) {
+            ctx->prealloc_y_last_decode_vector_staging != y_decode_vector_staging) {
             if (ctx->prealloc_y_need_sync) {
                 ggml_vk_sync_buffers(ctx, subctx);
             }
-            if (y_needs_k_padding) {
-                GGML_ASSERT(y_sz % 4 == 0);
-                // Zero B padding because clamping only A can produce 0 * Inf or NaN.
-                subctx->s->buffer->buf.fillBuffer(d_Y->buffer, 0, y_sz, 0);
-                ggml_vk_sync_buffers(ctx, subctx);
+            if (y_decode_vector_staging) {
                 const ggml_tensor y_staged_dst = make_y_staged_dst();
                 const uint32_t y_staged_dst_type_size = ggml_type_size(y_staged_dst.type);
                 ggml_vk_cpy_to_strided(
@@ -10458,27 +10428,27 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
             }
             ctx->prealloc_y_last_pipeline_used = to_fp16_vk_1.get();
             ctx->prealloc_y_last_tensor_used = src1;
-            ctx->prealloc_y_last_k_padded = y_needs_k_padding;
+            ctx->prealloc_y_last_decode_vector_staging = y_decode_vector_staging;
         }
     }
     if (quantize_y) {
         if (ctx->prealloc_y_last_pipeline_used != to_q8_1.get() ||
             ctx->prealloc_y_last_tensor_used != src1 ||
-            ctx->prealloc_y_last_k_padded) {
+            ctx->prealloc_y_last_decode_vector_staging) {
             if (ctx->prealloc_y_need_sync) {
                 ggml_vk_sync_buffers(ctx, subctx);
             }
             ggml_vk_quantize_q8_1(ctx, subctx, ggml_vk_subbuffer(ctx, d_Qy, qy_buf_offset), ggml_vk_subbuffer(ctx, d_Y, 0), y_ne);
             ctx->prealloc_y_last_pipeline_used = to_q8_1.get();
             ctx->prealloc_y_last_tensor_used = src1;
-            ctx->prealloc_y_last_k_padded = false;
+            ctx->prealloc_y_last_decode_vector_staging = false;
         }
     }
     ggml_vk_sync_buffers(ctx, subctx);
 
     uint32_t stride_batch_x = ne00*ne01;
-    uint32_t stride_b_y = y_needs_k_padding ? y_staged_row_stride : ne10;
-    uint32_t stride_batch_y = y_needs_k_padding ? y_staged_row_stride * ne11 : ne10*ne11;
+    uint32_t stride_b_y = y_decode_vector_staging ? y_staged_row_stride : ne10;
+    uint32_t stride_batch_y = y_decode_vector_staging ? y_staged_row_stride * padded_n : ne10*ne11;
 
     if (!ggml_vk_dim01_contiguous(src0) && !qx_needs_dequant) {
         stride_batch_x = src0->nb[0] / ggml_type_size(src0->type);
@@ -10495,13 +10465,13 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
         { d_D, d_buf_offset, d_sz }, { d_ids, ids_buf_offset, ids_sz }, expert_count_buf,
         ne01, ne21, ne10, ne10, stride_b_y, ne01,
         stride_batch_x, stride_batch_y, ne20*ne21,
-        n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, hoist_row_ids
+        n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, padded_n
     );  // NOLINT
 
     if (x_non_contig || qx_needs_dequant) {
         ctx->prealloc_x_need_sync = true;
     }
-    if (y_needs_reformat || quantize_y) {
+    if (y_non_contig || quantize_y) {
         ctx->prealloc_y_need_sync = true;
     }
     ctx->prealloc_split_k_need_sync = true;
@@ -10652,27 +10622,27 @@ static void ggml_vk_mul_mat_vec_id_q_f16(ggml_backend_vk_context * ctx, vk_conte
         GGML_ASSERT(y_sz == ggml_type_size(src1->type) * y_ne);
         if (ctx->prealloc_y_last_pipeline_used != to_fp16_vk_1.get() ||
             ctx->prealloc_y_last_tensor_used != src1 ||
-            ctx->prealloc_y_last_k_padded) {
+            ctx->prealloc_y_last_decode_vector_staging) {
             if (ctx->prealloc_y_need_sync) {
                 ggml_vk_sync_buffers(ctx, subctx);
             }
             ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, d_Qy, d_Y);
             ctx->prealloc_y_last_pipeline_used = to_fp16_vk_1.get();
             ctx->prealloc_y_last_tensor_used = src1;
-            ctx->prealloc_y_last_k_padded = false;
+            ctx->prealloc_y_last_decode_vector_staging = false;
         }
     }
     if (quantize_y) {
         if (ctx->prealloc_y_last_pipeline_used != to_q8_1.get() ||
             ctx->prealloc_y_last_tensor_used != src1 ||
-            ctx->prealloc_y_last_k_padded) {
+            ctx->prealloc_y_last_decode_vector_staging) {
             if (ctx->prealloc_y_need_sync) {
                 ggml_vk_sync_buffers(ctx, subctx);
             }
             ggml_vk_quantize_q8_1(ctx, subctx, d_Qy, d_Y, y_ne);
             ctx->prealloc_y_last_pipeline_used = to_q8_1.get();
             ctx->prealloc_y_last_tensor_used = src1;
-            ctx->prealloc_y_last_k_padded = false;
+            ctx->prealloc_y_last_decode_vector_staging = false;
         }
     }
 
@@ -15524,7 +15494,7 @@ static void ggml_vk_preallocate_buffers(ggml_backend_vk_context * ctx, vk_contex
         ctx->prealloc_y = ggml_vk_create_buffer_device(ctx->device, ctx->prealloc_size_y);
         ctx->prealloc_y_last_pipeline_used = nullptr;
         ctx->prealloc_y_last_tensor_used = nullptr;
-        ctx->prealloc_y_last_k_padded = false;
+        ctx->prealloc_y_last_decode_vector_staging = false;
     }
     if (ctx->prealloc_split_k == nullptr || (ctx->prealloc_size_split_k > 0 && ctx->prealloc_split_k->size < ctx->prealloc_size_split_k)) {
         VK_LOG_MEMORY("ggml_vk_preallocate_buffers(split_k_size: " << ctx->prealloc_size_split_k << ")");
@@ -16149,7 +16119,7 @@ static void ggml_vk_graph_cleanup(ggml_backend_vk_context * ctx) {
     VK_LOG_DEBUG("ggml_vk_graph_cleanup()");
     ctx->prealloc_y_last_pipeline_used = {};
     ctx->prealloc_y_last_tensor_used = nullptr;
-    ctx->prealloc_y_last_k_padded = false;
+    ctx->prealloc_y_last_decode_vector_staging = false;
 
     ctx->unsynced_nodes_written.clear();
     ctx->unsynced_nodes_read.clear();
@@ -16201,7 +16171,7 @@ static void ggml_vk_cleanup(ggml_backend_vk_context * ctx) {
 
     ctx->prealloc_y_last_pipeline_used = nullptr;
     ctx->prealloc_y_last_tensor_used = nullptr;
-    ctx->prealloc_y_last_k_padded = false;
+    ctx->prealloc_y_last_decode_vector_staging = false;
 
     ctx->prealloc_size_x = 0;
     ctx->prealloc_size_y = 0;
@@ -17399,7 +17369,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
 
     ctx->prealloc_y_last_pipeline_used = nullptr;
     ctx->prealloc_y_last_tensor_used = nullptr;
-    ctx->prealloc_y_last_k_padded = false;
+    ctx->prealloc_y_last_decode_vector_staging = false;
 
     if (ctx->prealloc_size_add_rms_partials) {
         ggml_vk_preallocate_buffers(ctx, nullptr);
@@ -17795,9 +17765,8 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
 }
 
 // Sort the graph for improved parallelism.
-static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * graph, struct ggml_backend_graph_optimize_params * params)
+static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * graph)
 {
-    GGML_UNUSED(params);
     VK_LOG_DEBUG("ggml_vk_graph_optimize(" << graph->n_nodes << " nodes)");
     ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context;
 
@@ -17805,32 +17774,20 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph *
         return;
     }
 
-    auto const &is_empty = [](const ggml_tensor * node) -> bool {
+    auto const &is_empty = [](ggml_tensor * node) -> bool {
         return node->op == GGML_OP_NONE || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE;
     };
 
-    auto const &is_src_of = [&is_empty](const ggml_tensor *dst, const ggml_tensor *src) -> bool {
-        auto const &base = [](const ggml_tensor * tensor) {
-            return tensor->view_src ? tensor->view_src : tensor;
-        };
+    auto const &is_src_of = [](const ggml_tensor *dst, const ggml_tensor *src) -> bool {
         for (uint32_t s = 0; s < GGML_MAX_SRC; ++s) {
             if (dst->src[s] == src) {
                 return true;
             }
-            if (is_empty(dst) || is_empty(src)) {
-                continue;
-            }
-            // A source view of dst may read storage written through a different view by src.
-            if (dst->src[s] && base(dst->src[s]) == base(src)) {
-                return true;
-            }
-            // Moving dst forward may overwrite storage still read through a view by src.
-            if (src->src[s] && base(dst) == base(src->src[s])) {
-                return true;
-            }
         }
         // implicit dependency if they view the same tensor
-        if (base(dst) == base(src)) {
+        const ggml_tensor *dst2 = dst->view_src ? dst->view_src : dst;
+        const ggml_tensor *src2 = src->view_src ? src->view_src : src;
+        if (dst2 == src2) {
             return true;
         }
         return false;

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

@@ -19,7 +19,6 @@
 #endif
 
 #include "types.glsl"
-#include "utils.glsl"
 
 // shape notation: [dim(N), ..., dim(0)] -- stride(dim(j)) >= stride(dim(i)) if i > j
 layout(binding = 0) readonly buffer A {
@@ -194,6 +193,14 @@ uint32_t       Br    = tid / BS_NPQ;
 uint32_t       Bc    = tid % BS_NPQ;
 const uint32_t BrpWg = WG_SIZE / BS_NPQ;
 
+// see init_fastdiv_values in ggml-vulkan.cpp
+uint fastdiv(uint n, uint mp, uint L) {
+    uint msbs, lsbs;
+    // msbs = mulhi(n, mp)
+    umulExtended(n, mp, msbs, lsbs);
+    return (msbs + n) >> L;
+}
+
 #ifdef COOPMAT2
 #define ACC_TYPE float16_t
 

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

@@ -15,7 +15,6 @@
 #endif
 
 #include "types.glsl"
-#include "utils.glsl"
 
 // shape notation: [dim(N), ..., dim(0)] -- stride(dim(j)) >= stride(dim(i)) if i > j
 layout(binding = 0) readonly buffer A {
@@ -179,6 +178,14 @@ uint32_t       Br    = tid / BS_NPQ;
 uint32_t       Bc    = tid % BS_NPQ;
 const uint32_t BrpWg = WG_SIZE / BS_NPQ;
 
+// see init_fastdiv_values in ggml-vulkan.cpp
+uint fastdiv(uint n, uint mp, uint L) {
+    uint msbs, lsbs;
+    // msbs = mulhi(n, mp)
+    umulExtended(n, mp, msbs, lsbs);
+    return (msbs + n) >> L;
+}
+
 void split_crs(uint32_t crs_idx, out uint32_t ic, out uint32_t kd, out uint32_t kh, out uint32_t kw) {
     const uint32_t KHKW = KH * KW;
     const uint32_t KDKHKW = KD * KHKW;

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

@@ -2,13 +2,7 @@
 
 #extension GL_EXT_control_flow_attributes : enable
 
-#ifdef USE_SUBGROUPS
-#extension GL_KHR_shader_subgroup_basic : enable
-#extension GL_KHR_shader_subgroup_arithmetic : enable
-#endif
-
 #include "types.glsl"
-#include "utils.glsl"
 
 layout (push_constant) uniform parameter
 {
@@ -17,10 +11,6 @@ layout (push_constant) uniform parameter
     uint32_t nb00;
     uint32_t nb01;
     uint32_t a_offset;
-    uint32_t n_experts;
-    uint32_t hoist_row_ids;
-    uint32_t ne00mp;
-    uint32_t ne00L;
 } p;
 
 #define BLOCK_SIZE 256
@@ -31,90 +21,16 @@ layout (binding = 0) readonly buffer A {uint data_a[];};
 layout (binding = 1) writeonly buffer D {uint data_d[];};
 
 shared uint vals[BLOCK_SIZE];
-shared uint offsets[BLOCK_SIZE];
-shared uint cursors[BLOCK_SIZE];
 
-// data_d layout when p.hoist_row_ids is set:
-//   [0,              n_experts)   per-expert row count
-//   [n_experts,    2*n_experts)   per-expert start offset into the row id region
-//   [2*n_experts]                 total row count
-//   [2*n_experts + 1,         )   row ids grouped by expert, packed as (i01 << 16) | (i00 & 0xffff)
-// Otherwise only data_d[expert_id] is written, holding that expert's row count.
 void main() {
     const uint expert_id = gl_WorkGroupID.x;
     const uint num_elements = p.ne00 * p.ne01;
     const uint tid = gl_LocalInvocationID.x;
 
-    if (p.hoist_row_ids != 0) {
-        if (tid < p.n_experts) {
-            vals[tid] = 0;
-        }
-        barrier();
-
-        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];
-            if (expert < p.n_experts) {
-                atomicAdd(vals[expert], 1);
-            }
-        }
-        barrier();
-
-#ifdef USE_SUBGROUPS
-        if (gl_SubgroupID == 0) {
-            // pad the trip count so the subgroup ops stay in uniform control flow
-            const uint n_experts_padded = (p.n_experts + gl_SubgroupSize - 1) & ~(gl_SubgroupSize - 1);
-            uint base = 0;
-            for (uint expert = gl_SubgroupInvocationID; expert < n_experts_padded; expert += gl_SubgroupSize) {
-                const bool in_range = expert < p.n_experts;
-                const uint count = in_range ? vals[expert] : 0;
-                const uint offset = base + subgroupExclusiveAdd(count);
-                if (in_range) {
-                    data_d[expert] = count;
-                    data_d[p.n_experts + expert] = offset;
-                    offsets[expert] = offset;
-                    cursors[expert] = 0;
-                }
-                base += subgroupAdd(count);
-            }
-            if (subgroupElect()) {
-                data_d[2 * p.n_experts] = base;
-            }
-        }
-#else
-        if (tid == 0) {
-            uint offset = 0;
-            for (uint expert = 0; expert < p.n_experts; ++expert) {
-                const uint count = vals[expert];
-                data_d[expert] = count;
-                data_d[p.n_experts + expert] = offset;
-                offsets[expert] = offset;
-                cursors[expert] = 0;
-                offset += count;
-            }
-            data_d[2 * p.n_experts] = offset;
-        }
-#endif
-        barrier();
-
-        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];
-            if (expert < p.n_experts) {
-                const uint row = atomicAdd(cursors[expert], 1);
-                const uint packed_row_id = (i01 << 16) | (i00 & 0xffffu);
-                data_d[2 * p.n_experts + 1 + offsets[expert] + row] = packed_row_id;
-            }
-        }
-        return;
-    }
-
     uint count = 0;
     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 i01 = idx / p.ne00;
+        const uint i00 = idx % p.ne00;
         const uint a = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00];
 
         count += uint(a == expert_id);

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

@@ -1,8 +1,6 @@
 #extension GL_EXT_shader_16bit_storage : require
 #extension GL_EXT_control_flow_attributes : require
 
-#include "utils.glsl"
-
 layout (push_constant) uniform parameter
 {
     uint ne;
@@ -34,6 +32,18 @@ uint get_idx() {
 uint get_aoffset() { return p.misalign_offsets >> 16; }
 uint get_doffset() { return p.misalign_offsets & 0xFFFF; }
 
+// see init_fastdiv_values in ggml-vulkan.cpp
+uint fastdiv(uint n, uint mp, uint L) {
+    uint msbs, lsbs;
+    // msbs = mulhi(n, mp)
+    umulExtended(n, mp, msbs, lsbs);
+    return (msbs + n) >> L;
+}
+
+uint fastdiv_L(uint packed, uint slot) {
+    return (packed >> (slot * 8)) & 0x3Fu;
+}
+
 uint src0_idx(uint idx) {
     const uint i03 = fastdiv(idx, p.ne0_012mp, fastdiv_L(p.ne0_Ls, 0));
     const uint i03_offset = i03 * p.ne02*p.ne01*p.ne00;

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

@@ -1,7 +1,5 @@
 #extension GL_EXT_shader_16bit_storage : require
 
-#include "utils.glsl"
-
 
 layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
 
@@ -41,3 +39,9 @@ uint get_aoffset() { return p.misalign_offsets >> 16; }
 uint get_boffset() { return (p.misalign_offsets >> 8) & 0xFF; }
 uint get_doffset() { return p.misalign_offsets & 0xFF; }
 
+// see init_fastdiv_values in ggml-vulkan.cpp
+uint fastdiv(uint n, uint mp, uint L) {
+    uint msbs, lsbs;
+    umulExtended(n, mp, msbs, lsbs);
+    return (msbs + n) >> L;
+}

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

@@ -88,8 +88,6 @@ layout (push_constant) uniform parameter
     uint nei1;
     uint nbi1;
     uint ne11;
-    uint n_experts;
-    uint hoist_row_ids;
 #else
     uint base_work_group_z;
     uint num_batches;
@@ -216,32 +214,28 @@ void main() {
     const uint loadstride_b = gl_WorkGroupSize.x * LOAD_VEC_B_EFF * LOAD_VEC_BATCH_B / BK;
 
 #ifdef MUL_MAT_ID
-    if (p.hoist_row_ids != 0) {
-        load_row_ids_hoisted(expert_idx, ic);
-    } else {
 #ifdef MUL_MAT_ID_USE_SUBGROUPS
-        if (bitCount(p.nei0) == 1) {
-            load_row_ids(expert_idx, true, ic);
-        } else {
-            load_row_ids(expert_idx, false, ic);
-        }
+    if (bitCount(p.nei0) == 1) {
+        load_row_ids(expert_idx, true, ic);
+    } else {
+        load_row_ids(expert_idx, false, ic);
+    }
 #else
-        _ne1 = 0;
-        for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) {
-            for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) {
-                if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) {
-                    if (_ne1 >= ic * BN) {
-                        row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1);
-                    }
-                    _ne1++;
+    _ne1 = 0;
+    for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) {
+        for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) {
+            if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) {
+                if (_ne1 >= ic * BN) {
+                    row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1);
                 }
+                _ne1++;
             }
         }
-
-        barrier();
-#endif
     }
 
+    barrier();
+#endif
+
     // Workgroup has no work
     if (ic * BN >= _ne1) return;
 #endif

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

@@ -56,8 +56,6 @@ layout (push_constant) uniform parameter
     uint nei1;
     uint nbi1;
     uint ne11;
-    uint n_experts;
-    uint hoist_row_ids;
 #else
     uint base_work_group_z;
     uint num_batches;
@@ -66,9 +64,9 @@ layout (push_constant) uniform parameter
     uint ne12;
     uint broadcast2;
     uint broadcast3;
+#endif
     // N dimension for the B matrix can be >= p.N
     uint padded_N;
-#endif
 } p;
 
 
@@ -227,23 +225,6 @@ void load_row_ids(uint expert_idx, bool nei0_is_pow2, uint ic) {
     }
     barrier();
 }
-
-void load_row_ids_hoisted(uint expert_idx, uint ic) {
-    _ne1 = uint(data_expert_count[expert_idx]);
-
-    const uint tile_begin = ic * BN;
-    const uint tile_count = tile_begin < _ne1 ? min(BN, _ne1 - tile_begin) : 0;
-    const uint expert_offset = uint(data_expert_count[p.n_experts + expert_idx]);
-    const uint row_ids_offset = 2 * p.n_experts + 1 + expert_offset + tile_begin;
-
-    for (uint i = gl_LocalInvocationIndex; i < tile_count; i += BLOCK_SIZE) {
-        const uint packed_row_id = uint(data_expert_count[row_ids_offset + i]);
-        const uint ii0 = packed_row_id & 0xffffu;
-        const uint ii1 = packed_row_id >> 16;
-        row_ids[i] = u16vec4(fastmod(ii0, p.ne11), ii1, ii0, 0);
-    }
-    barrier();
-}
 #endif
 
 void main() {
@@ -285,9 +266,7 @@ void main() {
     const uint ik = gl_WorkGroupID.x / blocks_m;
 
 #ifdef MUL_MAT_ID
-    if (p.hoist_row_ids != 0) {
-        load_row_ids_hoisted(expert_idx, ic);
-    } else if (bitCount(p.nei0) == 1) {
+    if (bitCount(p.nei0) == 1) {
         load_row_ids(expert_idx, true, ic);
     } else {
         load_row_ids(expert_idx, false, ic);
@@ -330,9 +309,7 @@ void main() {
     tensorLayoutNV<2> tensorLayoutA = createTensorLayoutNV(2);
     tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutAClamp = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
     tensorLayoutNV<2> tensorLayoutB = createTensorLayoutNV(2);
-#ifndef MUL_MAT_ID
     tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutBClamp = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
-#endif
     tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutD = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
 
 #if QUANT_K > 1
@@ -345,19 +322,12 @@ void main() {
 
     // Use end_k rather than p.K as the dimension because that's what
     // we need to bound check against when using split_k.
-    tensorLayoutA = setTensorLayoutDimensionNV(tensorLayoutA, p.M, end_k);
-#ifdef MUL_MAT_ID
-    // MUL_MAT_ID pads each B row to stride_b so partial K tiles read zeros without clamping.
-    tensorLayoutB = setTensorLayoutDimensionNV(tensorLayoutB, BN, p.stride_b);
-#else
     // Bounds check B against padded_N, but bounds check D against N.
+    tensorLayoutA = setTensorLayoutDimensionNV(tensorLayoutA, p.M, end_k);
     tensorLayoutB = setTensorLayoutDimensionNV(tensorLayoutB, p.padded_N, end_k);
-#endif
     tensorLayoutD = setTensorLayoutDimensionNV(tensorLayoutD, p.N, p.M);
     tensorLayoutAClamp = setTensorLayoutDimensionNV(tensorLayoutAClamp, p.M, end_k);
-#ifndef MUL_MAT_ID
     tensorLayoutBClamp = setTensorLayoutDimensionNV(tensorLayoutBClamp, p.padded_N, end_k);
-#endif
 
     tensorLayoutD = setTensorLayoutStrideNV(tensorLayoutD, p.stride_d, 1);
 
@@ -534,9 +504,7 @@ void main() {
 
         tensorLayoutB = setTensorLayoutStrideNV(tensorLayoutB, stride_b, 1);
 
-#ifndef MUL_MAT_ID
         tensorLayoutBClamp = setTensorLayoutStrideNV(tensorLayoutBClamp, stride_b, 1);
-#endif
 
         uint k_iters = (end_k - start_k + BK - 1) / BK;
 

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

@@ -71,19 +71,4 @@ void load_row_ids(uint expert_idx, bool nei0_is_pow2, uint ic) {
     barrier();
 }
 #endif // MUL_MAT_ID_USE_SUBGROUPS
-
-void load_row_ids_hoisted(uint expert_idx, uint ic) {
-    _ne1 = uint(data_expert_count[expert_idx]);
-
-    const uint tile_begin = ic * BN;
-    const uint tile_count = tile_begin < _ne1 ? min(BN, _ne1 - tile_begin) : 0;
-    const uint expert_offset = uint(data_expert_count[p.n_experts + expert_idx]);
-    const uint row_ids_offset = 2 * p.n_experts + 1 + expert_offset + tile_begin;
-
-    for (uint i = gl_LocalInvocationIndex; i < tile_count; i += BLOCK_SIZE) {
-        const uint packed_row_id = uint(data_expert_count[row_ids_offset + i]);
-        row_ids[i] = u16vec2(packed_row_id & 0xffffu, packed_row_id >> 16);
-    }
-    barrier();
-}
 #endif // MUL_MAT_ID

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

@@ -56,8 +56,6 @@ layout (push_constant) uniform parameter
     uint nei1;
     uint nbi1;
     uint ne11;
-    uint n_experts;
-    uint hoist_row_ids;
 #else
     uint base_work_group_z;
     uint num_batches;
@@ -159,32 +157,28 @@ void main() {
     const uint loadstride_b = BLOCK_SIZE * LOAD_VEC_B / BK;
 
 #ifdef MUL_MAT_ID
-    if (p.hoist_row_ids != 0) {
-        load_row_ids_hoisted(expert_idx, ic);
-    } else {
 #ifdef MUL_MAT_ID_USE_SUBGROUPS
-        if (bitCount(p.nei0) == 1) {
-            load_row_ids(expert_idx, true, ic);
-        } else {
-            load_row_ids(expert_idx, false, ic);
-        }
+    if (bitCount(p.nei0) == 1) {
+        load_row_ids(expert_idx, true, ic);
+    } else {
+        load_row_ids(expert_idx, false, ic);
+    }
 #else
-        _ne1 = 0;
-        for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) {
-            for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) {
-                if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) {
-                    if (_ne1 >= ic * BN) {
-                        row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1);
-                    }
-                    _ne1++;
+    _ne1 = 0;
+    for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) {
+        for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) {
+            if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) {
+                if (_ne1 >= ic * BN) {
+                    row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1);
                 }
+                _ne1++;
             }
         }
-
-        barrier();
-#endif
     }
 
+    barrier();
+#endif
+
     // Workgroup has no work
     if (ic * BN >= _ne1) return;
 #endif

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

@@ -1,6 +1,4 @@
 
-#include "utils.glsl"
-
 // vk_op_sum_rows_push_constants
 layout (push_constant) uniform parameter
 {
@@ -17,3 +15,11 @@ layout (push_constant) uniform parameter
 uint get_aoffset() { return p.misalign_offsets >> 16; }
 uint get_doffset() { return p.misalign_offsets & 0xFFFF; }
 
+// see init_fastdiv_values in ggml-vulkan.cpp
+uint fastdiv(uint n, uint mp, uint L) {
+    uint msbs, lsbs;
+    // msbs = mulhi(n, mp)
+    umulExtended(n, mp, msbs, lsbs);
+    return (msbs + n) >> L;
+}
+

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

@@ -9,26 +9,14 @@ uint fastmod(uint a, uint b) {
     return a % b;
 }
 
-// see init_fastdiv_values in ggml-vulkan.cpp
-uint fastdiv(uint n, uint mp, uint L) {
-    uint msbs, lsbs;
-    // msbs = mulhi(n, mp)
-    umulExtended(n, mp, msbs, lsbs);
-    return (msbs + n) >> L;
-}
-
-uint fastdiv_L(uint packed, uint slot) {
-    return (packed >> (slot * 8)) & 0x3Fu;
-}
-
-uint fastdiv_small(uint a, uint b) {
+uint fastdiv(uint a, uint b) {
     return (a < b) ? 0 : (a / b);
 }
 
 void get_indices(uint idx, out uint i00, out uint i01, out uint i02, out uint i03, uint ne00, uint ne01, uint ne02, uint ne03) {
-    i03 = fastdiv_small(idx, (ne02*ne01*ne00));
+    i03 = fastdiv(idx, (ne02*ne01*ne00));
     const uint i03_offset = i03 * ne02*ne01*ne00;
-    i02 = fastdiv_small((idx - i03_offset), (ne01*ne00));
+    i02 = fastdiv((idx - i03_offset), (ne01*ne00));
     const uint i02_offset = i02*ne01*ne00;
     i01 = (idx - i03_offset - i02_offset) / ne00;
     i00 = idx - i03_offset - i02_offset - i01*ne00;

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

@@ -1039,7 +1039,6 @@ void process_shaders() {
     string_to_spv("cumsum_multipass2_f32", "cumsum_multipass2.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
 
     string_to_spv("count_experts", "count_experts.comp", merge_maps(base_dict, {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}}));
-    string_to_spv("count_experts_subgroup", "count_experts.comp", merge_maps(base_dict, {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}, {"USE_SUBGROUPS", "1"}}));
 
     for (std::string dim_str : {"", "_3d"}) {
         for (bool bda : {false, true}) {

include/llama.h

@@ -214,10 +214,10 @@ extern "C" {
     LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode);
     LLAMA_API enum llama_load_mode llama_load_mode_from_str(const char * str);
 
-    enum llama_lazy_mode {
-        LLAMA_LAZY_MODE_OFF  = 0, // always read the whole tensor up front
-        LLAMA_LAZY_MODE_AUTO = 1, // lazy only for marked tensors larger than 4 GiB (requires mmap)
-        LLAMA_LAZY_MODE_ON   = 2, // read the rows of tensors marked by the arch on demand (requires mmap)
+    enum llama_tensor_read_lazy {
+        LLAMA_TENSOR_READ_LAZY_OFF  = 0, // always read the whole tensor up front
+        LLAMA_TENSOR_READ_LAZY_AUTO = 1, // lazy only for marked tensors larger than 4 GiB (requires mmap)
+        LLAMA_TENSOR_READ_LAZY_ON   = 2, // read the rows of tensors marked by the arch on demand (requires mmap)
     };
 
     enum llama_context_type {
@@ -321,7 +321,7 @@ extern "C" {
         enum llama_split_mode split_mode; // how to split the model across multiple GPUs
         enum llama_load_mode  load_mode;  // how to load the model
 
-        enum llama_lazy_mode lazy_mode; // on-demand reading of tensors marked by the arch
+        enum llama_tensor_read_lazy tensor_read_lazy; // on-demand reading of tensors marked by the arch
 
         // the GPU that is used for the entire model when split_mode is LLAMA_SPLIT_MODE_NONE
         int32_t main_gpu;

scripts/snapdragon/build.py

@@ -11,8 +11,6 @@ import platform
 import shutil
 import logging
 
-from sdk import validate_windows_sdks
-
 logger = logging.getLogger("build")
 
 
@@ -67,13 +65,6 @@ def main():
         logger.error(f"Error: Invalid target format '{args.target}'. Must be android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, or windows/wos.")
         sys.exit(1)
 
-    if target_type == "windows":
-        logger.info("Windows target selected. Forcing native compilation...")
-        args.no_docker = True
-        if platform.system() != "Windows":
-            logger.warning("Warning: Windows compilation is intended to run on Windows arm64 hosts.")
-        validate_windows_sdks()
-
     # Determine preset and check if it's debug
     preset = args.preset
     if preset:
@@ -129,6 +120,12 @@ def main():
 
     jobs = args.jobs if args.jobs else os.cpu_count() or 4
 
+    if target_type == "windows":
+        logger.info("Windows target selected. Forcing native compilation...")
+        args.no_docker = True
+        if platform.system() != "Windows":
+            logger.warning("Warning: Windows compilation is intended to run on Windows arm64 hosts.")
+
     if args.no_docker:
         # Native/local host build
         logger.info("Running native/local CMake build...")
@@ -261,6 +258,3 @@ if __name__ == "__main__":
     except KeyboardInterrupt:
         logger.info("\nInterrupted by user.")
         sys.exit(130)
-    except RuntimeError as err:
-        logger.error("Error: %s", err)
-        sys.exit(1)

scripts/snapdragon/sdk.py

@@ -1,62 +0,0 @@
-import os
-from pathlib import Path
-
-
-SDK_CONFIGS = (
-    {
-        "name": "Hexagon SDK",
-        "repo": "snapdragon-toolchain/hexagon-sdk",
-        "default_version": "6.6.0.0",
-        "parent_dir": "Hexagon_SDK",
-        "archive_prefix": "hexagon-sdk-v",
-        "markers": ("hexagon_sdk.json",),
-    },
-    {
-        "name": "OpenCL SDK",
-        "repo": "snapdragon-toolchain/opencl-sdk",
-        "default_version": "2.3.2",
-        "parent_dir": "OpenCL_SDK",
-        "archive_prefix": "adreno-opencl-sdk-v",
-        "markers": ("include/CL", "lib/OpenCL.lib"),
-    },
-)
-
-
-def is_valid_sdk(config, target_dir):
-    return target_dir.is_dir() and all((target_dir / marker).exists() for marker in config["markers"])
-
-
-def get_hexagon_tools_dir(hexagon_dir):
-    tools_parent = hexagon_dir / "tools" / "HEXAGON_Tools"
-    if not tools_parent.is_dir():
-        raise RuntimeError(f"Expected Hexagon tools directory in {tools_parent}")
-    tools_dirs = [path for path in tools_parent.iterdir() if path.is_dir()]
-    if len(tools_dirs) != 1:
-        raise RuntimeError(f"Expected one Hexagon tools directory in {tools_parent}")
-    return tools_dirs[0]
-
-
-def validate_windows_sdks():
-    hexagon_config, opencl_config = SDK_CONFIGS
-    hexagon_dir = os.environ.get("HEXAGON_SDK_ROOT")
-    tools_dir = os.environ.get("HEXAGON_TOOLS_ROOT")
-    opencl_dir = os.environ.get("OPENCL_SDK_ROOT")
-    missing = []
-
-    expected_tools_dir = None
-    if not hexagon_dir or not is_valid_sdk(hexagon_config, Path(hexagon_dir)):
-        missing.append("HEXAGON_SDK_ROOT")
-    else:
-        try:
-            expected_tools_dir = get_hexagon_tools_dir(Path(hexagon_dir))
-        except RuntimeError:
-            pass
-    if not tools_dir or not expected_tools_dir or Path(tools_dir) != expected_tools_dir:
-        missing.append("HEXAGON_TOOLS_ROOT")
-    if not opencl_dir or not is_valid_sdk(opencl_config, Path(opencl_dir)):
-        missing.append("OPENCL_SDK_ROOT")
-    if missing:
-        raise RuntimeError(
-            f"Missing or invalid Windows SDK paths: {', '.join(missing)}. "
-            "Run scripts/snapdragon/setup-sdk.py first."
-        )

scripts/snapdragon/setup-sdk.py

@@ -1,233 +0,0 @@
-#!/usr/bin/env python3
-#
-# Install Windows on Snapdragon SDKs for llama.cpp.
-#
-
-import sys
-import os
-import argparse
-import shutil
-import logging
-import json
-import hashlib
-import tarfile
-import tempfile
-from pathlib import Path
-from urllib.error import HTTPError, URLError
-from urllib.request import Request, urlopen
-
-from sdk import SDK_CONFIGS, get_hexagon_tools_dir, is_valid_sdk
-
-
-logger = logging.getLogger("setup_sdk")
-
-DEFAULT_SDK_BASE_DIR = r"C:\Qualcomm"
-
-
-def get_sdk_releases(config):
-    request = Request(
-        f"https://api.github.com/repos/{config['repo']}/releases?per_page=100",
-        headers={"Accept": "application/vnd.github+json", "User-Agent": "llama.cpp"},
-    )
-    try:
-        with urlopen(request, timeout=30) as response:
-            releases = json.load(response)
-    except (HTTPError, URLError, TimeoutError) as err:
-        raise RuntimeError(f"Cannot query {config['name']} releases: {err}") from err
-
-    result = []
-    for release in releases:
-        if release["draft"] or release["prerelease"]:
-            continue
-        version = release["tag_name"].removeprefix("v")
-        archive_name = f"{config['archive_prefix']}{version}-arm64-wos.tar.xz"
-        for asset in release["assets"]:
-            if asset["name"] != archive_name:
-                continue
-            result.append({
-                "version": version,
-                "name": asset["name"],
-                "url": asset["browser_download_url"],
-                "sha256": (asset.get("digest") or "").removeprefix("sha256:"),
-            })
-    return result
-
-
-def list_sdk_releases():
-    for config in SDK_CONFIGS:
-        logger.info("%s:", config["name"])
-        releases = get_sdk_releases(config)
-        if not releases:
-            logger.info("  no Windows on Snapdragon releases found")
-            continue
-        for release in releases:
-            logger.info("  %s: %s", release["version"], release["name"])
-
-
-def get_sdk_release(config, version):
-    version = version or config["default_version"]
-    version = version.removeprefix("v")
-    for release in get_sdk_releases(config):
-        if release["version"] == version:
-            if not release["sha256"]:
-                raise RuntimeError(f"{config['name']} {version} does not provide a SHA-256 digest")
-            return release
-    raise RuntimeError(
-        f"No Windows on Snapdragon release for {config['name']} {version}. "
-        "Run scripts/snapdragon/setup-sdk.py --list-sdk-releases to see available versions."
-    )
-
-
-def sha256sum(path):
-    digest = hashlib.sha256()
-    with open(path, "rb") as file:
-        for chunk in iter(lambda: file.read(1024 * 1024), b""):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
-def download_sdk(release, archive):
-    while True:
-        if archive.exists() and sha256sum(archive) == release["sha256"]:
-            logger.info("Using existing archive %s", archive)
-            return
-
-        offset = archive.stat().st_size if archive.exists() else 0
-        headers = {"User-Agent": "llama.cpp"}
-        if offset:
-            headers["Range"] = f"bytes={offset}-"
-            logger.info("Resuming download of %s at %d MiB", release["name"], offset // (1024 * 1024))
-        else:
-            logger.info("Downloading %s", release["name"])
-
-        try:
-            with urlopen(Request(release["url"], headers=headers), timeout=30) as response:
-                mode = "ab" if offset and response.status == 206 else "wb"
-                with open(archive, mode) as file:
-                    shutil.copyfileobj(response, file)
-        except HTTPError as err:
-            if err.code != 416:
-                raise RuntimeError(f"Cannot download {release['name']}: {err}") from err
-            archive.unlink(missing_ok=True)
-            continue
-        except (URLError, TimeoutError) as err:
-            raise RuntimeError(f"Cannot download {release['name']}: {err}") from err
-
-        if sha256sum(archive) == release["sha256"]:
-            return
-        raise RuntimeError(f"SHA-256 mismatch for {archive}. Re-run the command to resume the download.")
-
-
-def extract_sdk(config, archive, target_dir):
-    if not hasattr(tarfile, "data_filter"):
-        raise RuntimeError("SDK extraction requires Python 3.10.12 or later")
-
-    with tempfile.TemporaryDirectory(prefix=f".{target_dir.name}.tmp-", dir=target_dir.parent) as staging_path:
-        staging_dir = Path(staging_path)
-        with tarfile.open(archive, "r:xz") as tar:
-            tar.extractall(staging_dir, filter=tarfile.data_filter)
-
-        candidates = [staging_dir] + [path for path in staging_dir.iterdir() if path.is_dir()]
-        extracted_dirs = [path for path in candidates if is_valid_sdk(config, path)]
-        if len(extracted_dirs) != 1:
-            raise RuntimeError(f"{config['name']} archive does not contain the expected files")
-        extracted_dir = extracted_dirs[0]
-
-        backup_dir = None
-        if target_dir.exists():
-            backup_dir = target_dir.parent / f".{target_dir.name}.backup"
-            if backup_dir.exists():
-                raise RuntimeError(f"Cannot replace {target_dir}: backup directory {backup_dir} already exists")
-            target_dir.replace(backup_dir)
-        try:
-            extracted_dir.replace(target_dir)
-        except Exception:
-            if backup_dir:
-                backup_dir.replace(target_dir)
-            raise
-        if backup_dir:
-            shutil.rmtree(backup_dir)
-
-
-def install_sdk(config, version, base_dir, force):
-    version = (version or config["default_version"]).removeprefix("v")
-    target_dir = base_dir / config["parent_dir"] / version
-    if is_valid_sdk(config, target_dir) and not force:
-        logger.info("Using existing %s at %s", config["name"], target_dir)
-        return target_dir
-
-    release = get_sdk_release(config, version)
-    target_dir.parent.mkdir(parents=True, exist_ok=True)
-    archive = target_dir.parent / release["name"]
-    download_sdk(release, archive)
-    logger.info("Extracting %s to %s", config["name"], target_dir)
-    extract_sdk(config, archive, target_dir)
-    archive.unlink(missing_ok=True)
-    return target_dir
-
-
-def set_user_environment(values):
-    if os.name != "nt":
-        raise RuntimeError("SDK setup must run on Windows")
-
-    import winreg
-
-    with winreg.CreateKey(winreg.HKEY_CURRENT_USER, "Environment") as key:
-        for name, value in values.items():
-            winreg.SetValueEx(key, name, 0, winreg.REG_SZ, str(value))
-            os.environ[name] = str(value)
-
-    import ctypes
-
-    result = ctypes.c_ulong()
-    ctypes.windll.user32.SendMessageTimeoutW(0xffff, 0x001a, 0, "Environment", 0x0002, 5000, ctypes.byref(result))
-
-
-def setup_sdks(args):
-    base_dir = Path(args.sdk_base_dir).expanduser().resolve()
-    hexagon_config, opencl_config = SDK_CONFIGS
-    environment = {}
-
-    if args.hexagon is not None:
-        hexagon_dir = install_sdk(hexagon_config, args.hexagon, base_dir, args.force)
-        environment["HEXAGON_SDK_ROOT"] = hexagon_dir
-        environment["HEXAGON_TOOLS_ROOT"] = get_hexagon_tools_dir(hexagon_dir)
-    if args.opencl is not None:
-        opencl_dir = install_sdk(opencl_config, args.opencl, base_dir, args.force)
-        environment["OPENCL_SDK_ROOT"] = opencl_dir
-
-    set_user_environment(environment)
-    logger.info("SDK environment variables were updated. Start a new terminal before building.")
-
-
-def main():
-    logging.basicConfig(level=logging.INFO, format="%(message)s")
-    parser = argparse.ArgumentParser(description="Install Windows on Snapdragon SDKs for llama.cpp.")
-    parser.add_argument("--list-sdk-releases", action="store_true", help="List available Windows on Snapdragon SDK releases")
-    parser.add_argument("--sdk-base-dir", default=DEFAULT_SDK_BASE_DIR, help=r"SDK installation directory (default: C:\Qualcomm)")
-    parser.add_argument("--hexagon", nargs="?", const=SDK_CONFIGS[0]["default_version"], metavar="VERSION", help="Install the Hexagon SDK, optionally selecting a version")
-    parser.add_argument("--opencl", nargs="?", const=SDK_CONFIGS[1]["default_version"], metavar="VERSION", help="Install the OpenCL SDK, optionally selecting a version")
-    parser.add_argument("--force", action="store_true", help="Reinstall selected SDKs even when they already exist")
-    args = parser.parse_args()
-
-    if args.list_sdk_releases:
-        if args.sdk_base_dir != DEFAULT_SDK_BASE_DIR or args.hexagon is not None or args.opencl is not None or args.force:
-            parser.error("Installation options cannot be combined with --list-sdk-releases")
-        list_sdk_releases()
-        return
-    if args.hexagon is None and args.opencl is None:
-        parser.error("Select at least one SDK with --hexagon or --opencl")
-    if os.name != "nt":
-        parser.error("SDK setup must run on Windows")
-    setup_sdks(args)
-
-
-if __name__ == "__main__":
-    try:
-        main()
-    except KeyboardInterrupt:
-        logger.info("\nInterrupted by user.")
-        sys.exit(130)
-    except RuntimeError as err:
-        logger.error("Error: %s", err)
-        sys.exit(1)

scripts/sync_vendor.py

@@ -5,7 +5,7 @@ import os
 import sys
 import subprocess
 
-HTTPLIB_VERSION = "refs/tags/v0.54.1"
+HTTPLIB_VERSION = "refs/tags/v0.53.1"
 
 # used by examples/gguf-hash, these repos have no release tag, so we pin a commit
 XXHASH_COMMIT      = "9f465f1ea932d6ad9a26cd77496311ffa544cd68"

src/llama-arch.cpp

@@ -1100,6 +1100,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) {
     switch (arch) {
         case LLM_ARCH_QWEN35:
         case LLM_ARCH_QWEN35MOE:
+        case LLM_ARCH_QWEN4EXP:
         case LLM_ARCH_DEEPSEEK4:
         case LLM_ARCH_NEMOTRON_H:
         case LLM_ARCH_NEMOTRON_H_MOE:
@@ -1141,6 +1142,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
         case LLM_ARCH_KIMI_LINEAR:
         case LLM_ARCH_BAILINGMOE3:
         case LLM_ARCH_KIMI_K3:
+        case LLM_ARCH_QWEN4EXP:
         case LLM_ARCH_QWEN3TTS:
             return false;
         default:

src/llama-context.cpp

@@ -231,10 +231,10 @@ llama_context::llama_context(
 
     cparams.fused_gdn_ar = true;
     cparams.fused_gdn_ch = true;
-    cparams.auto_fgdn    = false;
+    cparams.auto_fgdn    = true;
 
-    cparams.fused_lid = true;
-    cparams.auto_flid = false;
+    cparams.fused_lid    = true;
+    cparams.auto_flid    = true;
 
     cparams.fused_dsv4_hc_pre  = true;
     cparams.fused_dsv4_hc_comb = true;

src/llama-kv-cache.cpp

@@ -2034,7 +2034,7 @@ void llm_graph_input_k_shift::set_input(const llama_ubatch * ubatch) {
         kv_self->set_input_k_shift(k_shift);
     }
 
-    if (k_rot && k_rot->buffer) {
+    if (k_rot) {
         kv_self->set_input_k_rot(k_rot);
     }
 }

src/llama-memory-hybrid-idx.cpp

@@ -9,7 +9,9 @@
 #include <cassert>
 #include <cmath>
 #include <iterator>
+#include <limits>
 #include <stdexcept>
+#include <tuple>
 
 //
 // llama_memory_hybrid_idx
@@ -47,6 +49,7 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx(
     hparams_idx(model.hparams),
     mem_idx(filter_idx == nullptr ? nullptr : [&] {
         // MQA with a single key head of indexer_head_size, as llama_kv_cache_dsa shapes its own
+        hparams_idx.rope_type = LLAMA_ROPE_TYPE_NONE;
         std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1);
         hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size;
 
@@ -137,6 +140,8 @@ void llama_memory_hybrid_idx::clear(bool data) {
     if (mem_idx) {
         mem_idx->clear(data);
     }
+
+    qsa_histories.clear();
 }
 
 bool llama_memory_hybrid_idx::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
@@ -149,15 +154,96 @@ bool llama_memory_hybrid_idx::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_po
         mem_idx->seq_rm(seq_id, p0, p1);
     }
 
-    return get_mem_attn()->seq_rm(seq_id, p0, p1);
+    const bool res = get_mem_attn()->seq_rm(seq_id, p0, p1);
+    if (!res) {
+        return false;
+    }
+
+    auto remove = [&](qsa_history & history) {
+        history.erase(std::remove_if(history.begin(), history.end(), [&](const qsa_token & token) {
+            return (p0 < 0 || token.pos[0] >= p0) && (p1 < 0 || token.pos[0] < p1);
+        }), history.end());
+    };
+
+    if (seq_id < 0) {
+        for (auto & item : qsa_histories) {
+            remove(item.second);
+        }
+    } else {
+        auto it = qsa_histories.find(seq_id);
+        if (it != qsa_histories.end()) {
+            remove(it->second);
+            if (it->second.empty()) {
+                qsa_histories.erase(it);
+            }
+        }
+    }
+
+    return true;
 }
 
 void llama_memory_hybrid_idx::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
+    if (seq_id_src == seq_id_dst) {
+        return;
+    }
+
+    qsa_history copied;
+    const auto & cells_src = get_mem_attn()->get_cells(seq_id_src);
+    const auto & cells_dst = get_mem_attn()->get_cells(seq_id_dst);
+    const bool replace = &cells_src != &cells_dst;
+    const auto src = qsa_histories.find(seq_id_src);
+    if (src != qsa_histories.end()) {
+        using pos_key = std::tuple<llama_pos, llama_pos, llama_pos>;
+        std::map<pos_key, std::vector<bool>> cells_by_pos;
+        for (uint32_t cell = 0; cell < cells_src.size(); ++cell) {
+            if (cells_src.is_empty(cell) || !cells_src.seq_has(cell, seq_id_src)) {
+                continue;
+            }
+
+            const llama_pos pos = cells_src.pos_get(cell);
+            if ((p0 >= 0 && pos < p0) || (p1 >= 0 && pos >= p1)) {
+                continue;
+            }
+
+            const auto & ext = cells_src.ext_get(cell);
+            cells_by_pos[{ pos, ext.y, ext.x }].push_back(!replace && cells_src.seq_has(cell, seq_id_dst));
+        }
+
+        std::map<pos_key, size_t> next_cell;
+        for (const auto & token : src->second) {
+            if ((p0 >= 0 && token.pos[0] < p0) || (p1 >= 0 && token.pos[0] >= p1)) {
+                continue;
+            }
+
+            const pos_key key = { token.pos[0], token.pos[1], token.pos[2] };
+            auto cells = cells_by_pos.find(key);
+            if (cells == cells_by_pos.end()) {
+                continue;
+            }
+
+            size_t & index = next_cell[key];
+            if (index < cells->second.size() && !cells->second[index++]) {
+                copied.push_back(token);
+            }
+        }
+    }
+
     llama_memory_hybrid::seq_cp(seq_id_src, seq_id_dst, p0, p1);
 
     if (mem_idx) {
         mem_idx->seq_cp(seq_id_src, seq_id_dst, p0, p1);
     }
+
+    if (replace) {
+        if (copied.empty()) {
+            qsa_histories.erase(seq_id_dst);
+        } else {
+            qsa_histories[seq_id_dst] = std::move(copied);
+        }
+    } else if (!copied.empty()) {
+        auto & dst = qsa_histories[seq_id_dst];
+        dst.insert(dst.end(), copied.begin(), copied.end());
+    }
 }
 
 void llama_memory_hybrid_idx::seq_keep(llama_seq_id seq_id) {
@@ -166,6 +252,13 @@ void llama_memory_hybrid_idx::seq_keep(llama_seq_id seq_id) {
     if (mem_idx) {
         mem_idx->seq_keep(seq_id);
     }
+
+    auto it = qsa_histories.find(seq_id);
+    qsa_history keep = it == qsa_histories.end() ? qsa_history{} : std::move(it->second);
+    qsa_histories.clear();
+    if (!keep.empty()) {
+        qsa_histories.emplace(seq_id, std::move(keep));
+    }
 }
 
 void llama_memory_hybrid_idx::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
@@ -174,6 +267,15 @@ void llama_memory_hybrid_idx::seq_add(llama_seq_id seq_id, llama_pos p0, llama_p
     if (mem_idx) {
         mem_idx->seq_add(seq_id, p0, p1, shift);
     }
+
+    auto it = qsa_histories.find(seq_id);
+    if (it != qsa_histories.end()) {
+        for (auto & token : it->second) {
+            if ((p0 < 0 || token.pos[0] >= p0) && (p1 < 0 || token.pos[0] < p1)) {
+                token.pos[0] += shift;
+            }
+        }
+    }
 }
 
 void llama_memory_hybrid_idx::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
@@ -182,6 +284,15 @@ void llama_memory_hybrid_idx::seq_div(llama_seq_id seq_id, llama_pos p0, llama_p
     if (mem_idx) {
         mem_idx->seq_div(seq_id, p0, p1, d);
     }
+
+    auto it = qsa_histories.find(seq_id);
+    if (it != qsa_histories.end()) {
+        for (auto & token : it->second) {
+            if ((p0 < 0 || token.pos[0] >= p0) && (p1 < 0 || token.pos[0] < p1)) {
+                token.pos[0] /= d;
+            }
+        }
+    }
 }
 
 std::map<ggml_backend_buffer_type_t, size_t> llama_memory_hybrid_idx::memory_breakdown() const {
@@ -205,8 +316,28 @@ void llama_memory_hybrid_idx::state_write(llama_io_write_i & io, llama_seq_id se
         if (mem_idx) {
             mem_idx->state_write(io, seq_id, flags);
         }
-    }
 
+        uint32_t n_histories = 0;
+        if (seq_id < 0) {
+            n_histories = (uint32_t) qsa_histories.size();
+        } else if (qsa_histories.count(seq_id) != 0) {
+            n_histories = 1;
+        }
+        io.write(&n_histories, sizeof(n_histories));
+
+        for (const auto & item : qsa_histories) {
+            if (seq_id >= 0 && item.first != seq_id) {
+                continue;
+            }
+
+            io.write(&item.first, sizeof(item.first));
+            const uint64_t n_tokens = item.second.size();
+            io.write(&n_tokens, sizeof(n_tokens));
+            for (const auto & token : item.second) {
+                io.write(token.pos.data(), sizeof(token.pos));
+            }
+        }
+    }
 }
 
 void llama_memory_hybrid_idx::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
@@ -230,6 +361,34 @@ void llama_memory_hybrid_idx::state_read(llama_io_read_i & io, llama_seq_id seq_
             if (mem_idx) {
                 mem_idx->state_read_sinfo(io, seq_id, flags, nullptr, &sinfos_attn);
             }
+
+            uint32_t n_histories;
+            io.read(&n_histories, sizeof(n_histories));
+            if (n_histories > LLAMA_MAX_SEQ) {
+                throw std::runtime_error("invalid QSA history count");
+            }
+
+            if (seq_id < 0) {
+                qsa_histories.clear();
+            } else {
+                qsa_histories.erase(seq_id);
+            }
+
+            for (uint32_t ih = 0; ih < n_histories; ++ih) {
+                llama_seq_id stored_seq;
+                uint64_t n_tokens;
+                io.read(&stored_seq, sizeof(stored_seq));
+                io.read(&n_tokens, sizeof(n_tokens));
+                if (stored_seq < 0 || stored_seq >= LLAMA_MAX_SEQ || n_tokens > get_mem_attn()->get_size()) {
+                    throw std::runtime_error("invalid QSA history");
+                }
+
+                auto & history = qsa_histories[seq_id < 0 ? stored_seq : seq_id];
+                history.resize(n_tokens);
+                for (auto & token : history) {
+                    io.read(token.pos.data(), sizeof(token.pos));
+                }
+            }
         }
 
     } catch (...) {
@@ -255,12 +414,156 @@ void llama_memory_hybrid_idx::state_drop(llama_seq_id seq_id) {
     if (mem_idx) {
         mem_idx->seq_rm(seq_id, -1, -1);
     }
+
+    qsa_histories.erase(seq_id);
 }
 
 llama_kv_cache * llama_memory_hybrid_idx::get_mem_idx() const {
     return mem_idx.get();
 }
 
+void llama_memory_hybrid_idx::commit_qsa_tokens(const llama_ubatch & ubatch) {
+    if (!mem_idx) {
+        return;
+    }
+
+    for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
+        qsa_token token = {};
+        if (ubatch.token) {
+            token.pos = { ubatch.pos[i], ubatch.pos[i], ubatch.pos[i], 0 };
+        } else {
+            for (uint32_t ip = 0; ip < token.pos.size(); ++ip) {
+                token.pos[ip] = ip < ubatch.n_pos ? ubatch.pos[i + ip*ubatch.n_tokens] : ubatch.pos[i];
+            }
+        }
+
+        for (int32_t is = 0; is < ubatch.n_seq_id[i]; ++is) {
+            qsa_histories[ubatch.seq_id[i][is]].push_back(token);
+        }
+    }
+}
+
+void llama_memory_hybrid_idx::set_input_qsa(
+        ggml_tensor *       block_cells,
+        ggml_tensor *       block_pos,
+        ggml_tensor *       block_mask,
+        ggml_tensor *       selected,
+        const ggml_tensor * kq_mask,
+        const llama_ubatch * ubatch,
+        uint32_t            ratio,
+        uint32_t            block_topk) const {
+    GGML_ASSERT(ggml_backend_buffer_is_host(block_cells->buffer));
+    GGML_ASSERT(ggml_backend_buffer_is_host(block_pos->buffer));
+    GGML_ASSERT(ggml_backend_buffer_is_host(block_mask->buffer));
+    GGML_ASSERT(ggml_backend_buffer_is_host(selected->buffer));
+    GGML_ASSERT(ggml_backend_buffer_is_host(kq_mask->buffer));
+
+    const int64_t n_blocks = block_cells->ne[1];
+    const int64_t n_tokens = ubatch->n_tokens;
+    const int64_t n_pos    = block_pos->ne[1];
+    const int64_t n_kv     = selected->ne[0];
+
+    GGML_ASSERT(block_cells->type == GGML_TYPE_I32);
+    GGML_ASSERT(block_pos->type   == GGML_TYPE_I32);
+    GGML_ASSERT(block_mask->type  == GGML_TYPE_F32);
+    GGML_ASSERT(selected->type    == GGML_TYPE_F32);
+    GGML_ASSERT(block_cells->ne[0] == ratio && block_cells->ne[2] == n_tokens);
+    GGML_ASSERT(block_pos->ne[0] == n_blocks && block_pos->ne[2] == n_tokens);
+    GGML_ASSERT(block_mask->ne[0] == n_blocks && block_mask->ne[1] == n_tokens);
+    GGML_ASSERT(selected->ne[1] == n_tokens);
+    GGML_ASSERT(kq_mask->ne[0] == n_kv);
+
+    int32_t * cell_data     = (int32_t *) block_cells->data;
+    int32_t * pos_data      = (int32_t *) block_pos->data;
+    float *   mask_data     = (float *) block_mask->data;
+    float *   selected_data = (float *) selected->data;
+
+    std::fill(cell_data, cell_data + ggml_nelements(block_cells), 0);
+    std::fill(pos_data, pos_data + ggml_nelements(block_pos), 0);
+    std::fill(mask_data, mask_data + ggml_nelements(block_mask), -INFINITY);
+    std::fill(selected_data, selected_data + ggml_nelements(selected), 0.0f);
+
+    auto mask_visible = [&](int64_t query, uint32_t cell) {
+        const int64_t index = query*n_kv + cell;
+        if (kq_mask->type == GGML_TYPE_F16) {
+            return std::isfinite(ggml_fp16_to_fp32(((const ggml_fp16_t *) kq_mask->data)[index]));
+        }
+        return std::isfinite(((const float *) kq_mask->data)[index]);
+    };
+
+    for (int64_t iq = 0; iq < n_tokens; ++iq) {
+        const llama_seq_id seq_id = ubatch->seq_id[iq][0];
+        const auto found = qsa_histories.find(seq_id);
+        if (found == qsa_histories.end()) {
+            continue;
+        }
+
+        const auto & cells = get_mem_attn()->get_cells(seq_id);
+        using pos_key = std::tuple<llama_pos, llama_pos, llama_pos>;
+        std::map<pos_key, std::vector<uint32_t>> cells_by_pos;
+        for (uint32_t cell = 0; cell < cells.size() && cell < (uint32_t) n_kv; ++cell) {
+            if (cells.is_empty(cell) || !cells.seq_has(cell, seq_id)) {
+                continue;
+            }
+
+            const auto & ext = cells.ext_get(cell);
+            cells_by_pos[{ cells.pos_get(cell), ext.y, ext.x }].push_back(cell);
+        }
+
+        std::map<pos_key, size_t> next_cell;
+        std::vector<std::pair<const qsa_token *, uint32_t>> visible;
+        for (const auto & token : found->second) {
+            const pos_key key = { token.pos[0], token.pos[1], token.pos[2] };
+            auto cells = cells_by_pos.find(key);
+            if (cells == cells_by_pos.end()) {
+                continue;
+            }
+
+            size_t & index = next_cell[key];
+            if (index >= cells->second.size()) {
+                continue;
+            }
+
+            const uint32_t cell = cells->second[index++];
+            if (mask_visible(iq, cell)) {
+                visible.emplace_back(&token, cell);
+            }
+        }
+
+        const size_t n_complete = visible.size()/ratio;
+        const size_t n_write = std::min<size_t>(n_complete, n_blocks);
+        std::vector<uint8_t> used_cells(n_kv, 0);
+        for (size_t ib = 0; ib < n_write; ++ib) {
+            mask_data[iq*n_blocks + ib] = 0.0f;
+            for (uint32_t ir = 0; ir < ratio; ++ir) {
+                const uint32_t cell = visible[ib*ratio + ir].second;
+                cell_data[(iq*n_blocks + ib)*ratio + ir] = cell;
+                used_cells[cell] = 1;
+            }
+            for (int64_t ip = 0; ip < n_pos; ++ip) {
+                pos_data[(iq*n_pos + ip)*n_blocks + ib] = visible[ib*ratio].first->pos[ip];
+            }
+        }
+
+        uint32_t fallback_cell = 0;
+        for (size_t ib = n_write; ib < (size_t) n_blocks; ++ib) {
+            for (uint32_t ir = 0; ir < ratio; ++ir) {
+                while (fallback_cell < used_cells.size() && used_cells[fallback_cell]) {
+                    ++fallback_cell;
+                }
+                GGML_ASSERT(fallback_cell < used_cells.size());
+                cell_data[(iq*n_blocks + ib)*ratio + ir] = fallback_cell;
+                used_cells[fallback_cell++] = 1;
+            }
+        }
+
+        const size_t selected_start = n_complete <= block_topk ? 0 : n_complete*ratio;
+        for (size_t iv = selected_start; iv < visible.size(); ++iv) {
+            selected_data[iq*n_kv + visible[iv].second] = 1.0f;
+        }
+    }
+}
+
 //
 // llama_memory_hybrid_idx_context
 //
@@ -295,7 +598,9 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(
                   llama_context * lctx,
                            bool   optimize) :
     llama_memory_hybrid_context(mem, lctx, optimize),
-    mem(mem) {}
+    mem(mem),
+    ctx_idx(mem->get_mem_idx() == nullptr ? nullptr :
+        mem->get_mem_idx()->init_update(lctx, optimize)) {}
 
 llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(
         llama_memory_hybrid_idx * mem,
@@ -307,7 +612,8 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(
     mem(mem),
     ns_ubatch(llama_memory_hybrid_idx_ns(sinfos_idx)),
     ctx_idx(mem->get_mem_idx() == nullptr ? nullptr :
-        new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), ubatches)) {}
+        new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), ubatches)),
+    has_ubatches(true) {}
 
 bool llama_memory_hybrid_idx_context::next() {
     if (ctx_idx) {
@@ -326,6 +632,10 @@ bool llama_memory_hybrid_idx_context::apply() {
         res = res & ctx_idx->apply();
     }
 
+    if (res && ctx_idx && has_ubatches) {
+        mem->commit_qsa_tokens(ctx_idx->get_ubatch());
+    }
+
     return res;
 }
 
@@ -340,126 +650,17 @@ uint32_t llama_memory_hybrid_idx_context::get_n_stream() const {
 }
 
 void llama_memory_hybrid_idx_context::set_input_qsa(
-        ggml_tensor * cell_blk,
-        ggml_tensor * blk_cells,
-        ggml_tensor * blk_pos,
-        ggml_tensor * bias,
+        ggml_tensor *       block_cells,
+        ggml_tensor *       block_pos,
+        ggml_tensor *       block_mask,
+        ggml_tensor *       selected,
+        const ggml_tensor * kq_mask,
         const llama_ubatch * ubatch,
-        uint32_t ratio,
-        bool blk_bias) const {
+        uint32_t            ratio,
+        uint32_t            block_topk) const {
     GGML_ASSERT(ratio > 0);
     GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr);
 
-    GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer));
-
-    const int64_t n_kv     = cell_blk->ne[0];
-    const int64_t n_ns     = cell_blk->ne[1];        // streams in this ubatch
-    const int64_t n_blocks = blk_pos->ne[0]/(4*n_ns);
-    const int64_t n_tokens = ubatch->n_tokens;
-    const int64_t r        = ratio;
-
-    GGML_ASSERT(n_tokens % n_ns == 0);
-    const int64_t n_tps = n_tokens/n_ns;             // tokens per stream
-
-    int32_t * dst_cell_blk  = (int32_t *) cell_blk->data;
-    int32_t * dst_blk_cells = (int32_t *) blk_cells->data;
-    int32_t * dst_blk_pos   = (int32_t *) blk_pos->data;
-    float   * dst_bias      = (float   *) bias->data;
-
-    // block b covers [b*ratio, (b+1)*ratio), so its first token is at b*ratio
-    // all mrope sections carry it: exact for text, approximate for images
-    for (int64_t sec = 0; sec < 4; ++sec) {
-        for (int64_t s = 0; s < n_ns; ++s) {
-            for (int64_t b = 0; b < n_blocks; ++b) {
-                dst_blk_pos[sec*(n_blocks*n_ns) + s*n_blocks + b] = (int32_t) (b*r);
-            }
-        }
-    }
-
-    // one pass per stream: cell j is a different token in each, so no mapping is shared
-    std::vector<int32_t> blk_of(n_kv);
-    std::vector<int32_t> filled(n_blocks);
-
-    for (int64_t s = 0; s < n_ns; ++s) {
-        // ubatch index s*n_tps belongs to this stream; ask which cells array it uses
-        const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0];
-        const auto & cells = mem->get_mem_idx()->get_cells(seq_of_stream);
-
-        int32_t * cur_cell_blk  = dst_cell_blk  + s*n_kv;
-        int32_t * cur_blk_cells = dst_blk_cells + s*(r*n_blocks);
-
-        // an incomplete block cannot be pooled; the bias below forces those tail cells in
-        // -1 means no usable block, and block 0 only keeps the gather in range
-        std::fill(blk_of.begin(),  blk_of.end(),  -1);
-        std::fill(filled.begin(),  filled.end(),   0);
-        std::fill(cur_blk_cells, cur_blk_cells + r*n_blocks, 0);
-
-        // a cell no block covers needs its own -inf, which a per-block bias cannot carry
-        // every cache path keeps the position below the cell window, so this stays false
-        bool oor = false;
-
-        for (int64_t j = 0; j < n_kv; ++j) {
-            if (cells.is_empty(j)) {
-                continue;
-            }
-
-            const llama_pos p = cells.pos_get(j);
-            const int64_t   b = p/r;
-
-            if (b >= n_blocks) {
-                oor = true;
-                continue;
-            }
-
-            blk_of[j] = (int32_t) b;
-            cur_blk_cells[b*r + (p%r)] = (int32_t) j;
-            filled[b]++;
-        }
-
-        GGML_ASSERT((!blk_bias || !oor) && "qsa: cell position runs past the cell window");
-
-        // per-block mode keeps an unpooled cell's real block, so the block's own -inf reaches it
-        // per-cell mode carries that -inf itself and only needs the gather in range
-        for (int64_t j = 0; j < n_kv; ++j) {
-            if (blk_of[j] >= 0 && filled[blk_of[j]] < r && !blk_bias) {
-                blk_of[j] = -1;
-            }
-            cur_cell_blk[j] = blk_of[j] < 0 ? 0 : blk_of[j];
-        }
-
-        for (int64_t ii = 0; ii < n_tps; ++ii) {
-            const int64_t      i      = s*n_tps + ii;
-            const llama_seq_id seq_id = ubatch->seq_id[i][0];
-            const llama_pos    q      = ubatch->pos[i];
-
-            // the tail is an incomplete block and is always visible, as in the reference
-            const llama_pos tail_start = (q + 1)/r*r;
-
-            if (blk_bias) {
-                // a block sits wholly inside or outside the tail, so one value covers it
-                // the caller adds the attention mask, which drops empty, foreign and future cells
-                float * cur_blk_bias = dst_bias + i*n_blocks;
-
-                for (int64_t b = 0; b < n_blocks; ++b) {
-                    // finite, so it can never meet a -inf and produce a nan
-                    cur_blk_bias[b] = b*r >= tail_start ? 1e9f : (filled[b] < r ? -INFINITY : 0.0f);
-                }
-
-                continue;
-            }
-
-            float * cur_bias = dst_bias + i*n_kv;
-
-            for (int64_t j = 0; j < n_kv; ++j) {
-                float v = -INFINITY;
-
-                if (!cells.is_empty(j) && cells.seq_has(j, seq_id) && cells.pos_get(j) <= q) {
-                    // finite, so it can never meet a -inf and produce a nan
-                    v = cells.pos_get(j) >= tail_start ? 1e9f : (blk_of[j] < 0 ? -INFINITY : 0.0f);
-                }
-
-                cur_bias[j] = v;
-            }
-        }
-    }
+    mem->set_input_qsa(block_cells, block_pos, block_mask, selected,
+            kq_mask, ubatch, ratio, block_topk);
 }

src/llama-memory-hybrid-idx.h

@@ -2,6 +2,8 @@
 
 #include "llama-memory-hybrid.h"
 
+#include <array>
+#include <map>
 #include <memory>
 #include <vector>
 
@@ -75,7 +77,20 @@ public:
 
     llama_kv_cache * get_mem_idx() const;   // nullptr when the model carries no indexer
 
+    void set_input_qsa(ggml_tensor * block_cells, ggml_tensor * block_pos,
+                       ggml_tensor * block_mask, ggml_tensor * selected,
+                       const ggml_tensor * kq_mask, const llama_ubatch * ubatch,
+                       uint32_t ratio, uint32_t block_topk) const;
+
+    void commit_qsa_tokens(const llama_ubatch & ubatch);
+
 private:
+    struct qsa_token {
+        std::array<llama_pos, 4> pos;
+    };
+
+    using qsa_history = std::vector<qsa_token>;
+
     // forget seq_id (all of it if seq_id < 0) in every cache at once, so a failed restore cannot leave the caches out of step
     // seq_id < 0 drops the whole context, as the caches themselves do on a failed restore
     void state_drop(llama_seq_id seq_id);
@@ -85,6 +100,8 @@ private:
     llama_hparams hparams_idx;
 
     const std::unique_ptr<llama_kv_cache> mem_idx;
+
+    std::map<llama_seq_id, qsa_history> qsa_histories;
 };
 
 class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context {
@@ -123,26 +140,20 @@ public:
     // llama_memory_hybrid_idx_context specific API
     //
 
-    // nullptr with no indexer, and for the update context, which builds no sparse graph
+    // nullptr with no indexer
     const llama_kv_cache_context * get_idx() const;
 
     // streams in the current slot info, the `ns` of get_k/get_v; 1 if unified
     uint32_t get_n_stream() const;
 
-    // block-compressed sparse attention (qwen4exp QSA) over the cells of the indexer cache.
-    // Blocks cut the position line, not the cell array, so no caller assumes a contiguous layout:
-    //   cell_blk  I32 [n_kv, ns]           block each cell belongs to
-    //   blk_cells I32 [ratio*n_blocks, ns] cells making up each block
-    //   blk_pos   I32 [4*n_blocks*ns]      mrope position rows of each block's first token
-    //   bias      F32 [n_kv, n_tokens/ns, ns] -inf where invisible, large where always visible
-    // blk_bias asks for the bias per block instead: [n_blocks, n_tokens/ns, ns]
-    // the caller then adds the attention mask, the only part of the bias that varies within a block
-    void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos,
-                       ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio,
-                       bool blk_bias) const;
+    // QSA blocks follow each sequence's token order, not physical cells or scalar positions.
+    void set_input_qsa(ggml_tensor * block_cells, ggml_tensor * block_pos,
+                       ggml_tensor * block_mask, ggml_tensor * selected,
+                       const ggml_tensor * kq_mask, const llama_ubatch * ubatch,
+                       uint32_t ratio, uint32_t block_topk) const;
 
 private:
-    const llama_memory_hybrid_idx * mem = nullptr;
+    llama_memory_hybrid_idx * mem = nullptr;
 
     // streams per ubatch, read from the slot infos before ctx_idx takes them
     // declared first, so it is initialised while sinfos_idx is still intact
@@ -151,6 +162,8 @@ private:
     // null unless the model has an indexer and this is a batch or full context
     const llama_memory_context_ptr ctx_idx;
 
+    const bool has_ubatches = false;
+
     // mirrors the base class's ubatch cursor, which is private there
     size_t i_cur = 0;
 };

src/llama-model-loader.cpp

@@ -1287,10 +1287,10 @@ struct ggml_tensor * llama_model_loader::create_tensor(
         return NULL;
     }
 
-    if ((flags & TENSOR_READ_LAZY) && use_mmap && lazy_mode != LLAMA_LAZY_MODE_OFF) {
+    if ((flags & TENSOR_READ_LAZY) && use_mmap && tensor_read_lazy != LLAMA_TENSOR_READ_LAZY_OFF) {
         // in auto mode, small tensors are cheap enough to keep resident
         constexpr size_t auto_lazy_min_size = 4ull * 1024 * 1024 * 1024;
-        if (lazy_mode == LLAMA_LAZY_MODE_ON || ggml_nbytes(cur) > auto_lazy_min_size) {
+        if (tensor_read_lazy == LLAMA_TENSOR_READ_LAZY_ON || ggml_nbytes(cur) > auto_lazy_min_size) {
             const auto & w = require_weight(tn.str().c_str());
             lazy_tensor_ranges[w.idx].emplace_back(w.offs, w.offs + ggml_nbytes(cur));
 

src/llama-model-loader.h

@@ -84,7 +84,7 @@ struct llama_model_loader {
     bool load_mtp;
 
     // set by the caller before the create_tensor() calls
-    enum llama_lazy_mode lazy_mode = LLAMA_LAZY_MODE_OFF;
+    enum llama_tensor_read_lazy tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_OFF;
 
     llama_files files;
     llama_ftype ftype;

src/llama-model.cpp

@@ -2681,7 +2681,7 @@ llama_model_params llama_model_default_params() {
         /*.n_gpu_layers                =*/ -1,
         /*.split_mode                  =*/ LLAMA_SPLIT_MODE_LAYER,
         /*.load_mode                   =*/ LLAMA_LOAD_MODE_AUTO,
-        /*.lazy_mode                   =*/ LLAMA_LAZY_MODE_AUTO,
+        /*.tensor_read_lazy            =*/ LLAMA_TENSOR_READ_LAZY_AUTO,
         /*.main_gpu                    =*/ 0,
         /*.tensor_split                =*/ nullptr,
         /*.progress_callback           =*/ nullptr,

src/llama.cpp

@@ -318,7 +318,7 @@ static std::pair<int, llama_model *> llama_model_load(struct gguf_context * meta
         llama_model_loader ml(metadata, set_tensor_data, set_tensor_data_ud, fname, splits, file, params.load_mode,
             params.check_tensors, params.no_alloc, params.load_mtp, params.kv_overrides, params.tensor_buft_overrides);
 
-        ml.lazy_mode = params.lazy_mode;
+        ml.tensor_read_lazy = params.tensor_read_lazy;
 
         ml.print_info();
         std::unique_ptr<llama_model> model_ptr(llama_model_create(ml, params));

src/models/models.h

@@ -2310,22 +2310,12 @@ struct llama_model_qwen4exp : public llama_model_base {
                             int * sections,
                             int   il);
 
-        // dense self-attention restricted to the cells that top_k names
-        ggml_tensor * build_attn_qsa(
-        llm_graph_input_attn_kv * inp,
-                    ggml_tensor * q_cur,
-                    ggml_tensor * k_cur,
-                    ggml_tensor * v_cur,
-                    ggml_tensor * top_k,
-                          float   kq_scale,
-                            int   il);
-
         // the QSA cache layout inputs do not depend on the layer, only on its compress ratio,
         // so the layers sharing a ratio share one input set
         std::map<uint32_t, llm_graph_input_qsa *> qsa_inps;
 
-        // QSA: token indices this layer's queries may attend to, or nullptr for dense
-        ggml_tensor * build_qsa_top_k(
+        // QSA mask for this layer, or nullptr for dense attention
+        ggml_tensor * build_qsa_mask(
   const llama_memory_hybrid_idx_context * mctx_hyb,
                     ggml_tensor * cur,
                     ggml_tensor * inp_pos,
@@ -2360,12 +2350,9 @@ struct llama_model_qwen4exp : public llama_model_base {
                         int64_t   channels,
                             int   il);
 
-        ggml_tensor * build_inp_ple(
-  const llama_memory_hybrid_idx_context * mctx_hyb);
-
         ggml_tensor * build_ple(
              llm_graph_input_rs * inp,
-                    ggml_tensor * emb,
+  const llama_memory_hybrid_idx_context * mctx_hyb,
                     ggml_tensor * hidden,
                             int   il);
 

src/models/qwen4exp.cpp

@@ -18,21 +18,29 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
     ml.get_key(LLM_KV_SSM_STATE_SIZE,     hparams.ssm_d_state);
     ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);
     ml.get_key(LLM_KV_SSM_GROUP_COUNT,    hparams.ssm_n_group);
-    GGML_ASSERT(hparams.ssm_d_conv  > 0 && hparams.ssm_d_inner > 0 && hparams.ssm_d_state > 0 &&
-                hparams.ssm_dt_rank > 0 && hparams.ssm_n_group > 0);
+    if (hparams.ssm_d_conv == 0 || hparams.ssm_d_inner == 0 || hparams.ssm_d_state == 0 ||
+            hparams.ssm_dt_rank == 0 || hparams.ssm_n_group == 0 ||
+            hparams.ssm_dt_rank % hparams.ssm_n_group != 0 ||
+            (uint64_t) hparams.ssm_d_state * hparams.ssm_dt_rank != hparams.ssm_d_inner) {
+        throw std::runtime_error("invalid Qwen4-Exp gated delta net dimensions");
+    }
 
     // HC; low_rank is qwen4exp-specific, DeepSeek-V4 leaves it absent (full rank)
     ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT,    hparams.dsv4_hc_mult);
     ml.get_key(LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank);
-    GGML_ASSERT(hparams.dsv4_hc_mult > 0 && hparams.hc_low_rank > 0);
+    if (hparams.n_embd == 0 || hparams.dsv4_hc_mult <= 1 || hparams.hc_low_rank == 0 ||
+            hparams.dsv4_hc_mult > UINT32_MAX/hparams.n_embd) {
+        throw std::runtime_error("invalid Qwen4-Exp hyper-connection dimensions");
+    }
     hparams.n_embd_out_impl = hparams.dsv4_hc_mult * hparams.n_embd;
 
     ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);
     ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);
     ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K,      hparams.indexer_top_k);
-    GGML_ASSERT(hparams.indexer_n_head > 0
-             && hparams.indexer_head_size > 0
-             && hparams.indexer_top_k > 0);
+    if (hparams.indexer_n_head == 0 || hparams.indexer_head_size == 0 || hparams.indexer_top_k == 0 ||
+            hparams.n_rot_full > hparams.indexer_head_size) {
+        throw std::runtime_error("invalid Qwen4-Exp sparse-attention dimensions");
+    }
     ml.get_key_or_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, hparams.n_layer_all, false);
 
     // PLE n-gram hash embeddings; if the key group is absent every field stays zero
@@ -44,9 +52,11 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
     if (n_ple > 0) {
         std::vector<uint32_t> ple_layers;
         ml.get_arr(LLM_KV_PLE_LAYERS, ple_layers);
-        GGML_ASSERT(n_ple == 1 && "qwen4exp supports only one PLE layer");
+        if (n_ple != 1 || ple_layers.size() != n_ple) {
+            throw std::runtime_error("Qwen4-Exp supports exactly one PLE layer");
+        }
         for (uint32_t il : ple_layers) {
-            if (il >= hparams.n_layer_all) {
+            if (il >= hparams.n_layer()) {
                 throw std::runtime_error(format("PLE layer %u is out of range", il));
             }
             hparams.is_ple_impl.set(il);
@@ -59,15 +69,30 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
         // optional: files written before this key fall back to the EOS token
         ml.get_key(LLM_KV_PLE_IMAGE_TOKEN_ID,  hparams.ple_image_token_id, false);
         ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.n_embd_per_layer);
-        GGML_ASSERT(hparams.ple_conv_kernel > 0 && hparams.n_embd_per_layer > 0);
+        if (hparams.ple_conv_kernel == 0 || hparams.n_embd_per_layer == 0) {
+            throw std::runtime_error("invalid Qwen4-Exp PLE dimensions");
+        }
 
-        hparams.ple_n_heads  = (hparams.ple_ngram_size - 1) * hparams.ple_heads_per_ngram;
-        hparams.ple_head_dim = hparams.n_embd_per_layer;
         if (hparams.ple_ngram_size < 2 || hparams.ple_ngram_size > LLAMA_MAX_PLE_NGRAM) {
             throw std::runtime_error(format("PLE n-gram size %u is out of range", hparams.ple_ngram_size));
         }
-        if (hparams.ple_n_heads == 0 || hparams.ple_n_heads > LLAMA_MAX_PLE_HEADS) {
-            throw std::runtime_error(format("PLE head count %u is out of range", hparams.ple_n_heads));
+
+        const uint64_t ple_n_heads = (uint64_t) (hparams.ple_ngram_size - 1) * hparams.ple_heads_per_ngram;
+        hparams.ple_head_dim = hparams.n_embd_per_layer;
+        if (ple_n_heads == 0 || ple_n_heads > LLAMA_MAX_PLE_HEADS) {
+            throw std::runtime_error(format("PLE head count %" PRIu64 " is out of range", ple_n_heads));
+        }
+        hparams.ple_n_heads = (uint32_t) ple_n_heads;
+
+        uint32_t n_multipliers = 0;
+        uint32_t n_offsets = 0;
+        uint32_t n_vocab_sizes = 0;
+        ml.get_arr_n(LLM_KV_PLE_LAYER_MULTIPLIERS, n_multipliers);
+        ml.get_arr_n(LLM_KV_PLE_HEAD_OFFSETS,      n_offsets);
+        ml.get_arr_n(LLM_KV_PLE_HEAD_VOCAB_SIZES,  n_vocab_sizes);
+        if (n_multipliers != hparams.ple_ngram_size ||
+                n_offsets != hparams.ple_n_heads || n_vocab_sizes != hparams.ple_n_heads) {
+            throw std::runtime_error("invalid Qwen4-Exp PLE metadata lengths");
         }
 
         ml.get_arr(LLM_KV_PLE_LAYER_MULTIPLIERS, hparams.ple_layer_multipliers);
@@ -93,12 +118,28 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
     if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) {
         uint32_t full_attn_interval = 4;
         ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false);
-        GGML_ASSERT(full_attn_interval > 0);
+        if (full_attn_interval == 0) {
+            throw std::runtime_error("invalid Qwen4-Exp full-attention interval");
+        }
         for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {
             hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0);
         }
     }
 
+    for (uint32_t il = 0; il < hparams.n_layer(); ++il) {
+        const uint32_t ratio = hparams.dsv4_compress_ratios[il];
+        if (hparams.is_recr(il)) {
+            if (ratio != 0) {
+                throw std::runtime_error(format("Qwen4-Exp recurrent layer %u has a QSA compression ratio", il));
+            }
+        } else if (ratio == 0 || hparams.indexer_top_k % ratio != 0) {
+            throw std::runtime_error(format("invalid Qwen4-Exp QSA compression ratio %u at layer %u", ratio, il));
+        }
+        if (hparams.is_ple(il) && !hparams.is_recr(il)) {
+            throw std::runtime_error(format("Qwen4-Exp PLE layer %u is not recurrent", il));
+        }
+    }
+
     switch (hparams.n_layer()) {
         case 48: type = LLM_TYPE_A3B; break;
         default: type = LLM_TYPE_UNKNOWN;
@@ -127,8 +168,15 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) {
     // flat [ple_head_dim, n_rows] gather target; n_rows is padded, so read it back
     if (hparams.ple_n_heads > 0) {
         const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").str();
-        const auto & ple_w = ml.require_weight(ple_name.c_str());
-        const int64_t ple_rows = ple_w.tensor->ne[1];
+        const auto * ple_w = ml.get_weight(ple_name.c_str());
+        int64_t ple_rows = 0;
+        for (uint32_t h = 0; h < hparams.ple_n_heads; ++h) {
+            ple_rows = std::max<int64_t>(ple_rows,
+                    (int64_t) hparams.ple_head_offsets[h] + hparams.ple_head_vocab_sizes[h]);
+        }
+        if (ple_w) {
+            ple_rows = ple_w->tensor->ne[1];
+        }
 
         // sanity check
         for (uint32_t h = 0; h < hparams.ple_n_heads; ++h) {
@@ -190,8 +238,9 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) {
         }
 
         if (hparams.is_ple(il)) {
-            layer.ple_key        = create_tensor(tn(LLM_TENSOR_PLE_KEY,        "weight", il), { n_embd, hc_dim }, 0);
-            layer.ple_value      = create_tensor(tn(LLM_TENSOR_PLE_VALUE,      "weight", il), { n_embd, n_embd }, 0);
+            const int64_t ple_dim = (int64_t) hparams.ple_head_dim * hparams.ple_n_heads;
+            layer.ple_key        = create_tensor(tn(LLM_TENSOR_PLE_KEY,        "weight", il), { ple_dim, hc_dim }, 0);
+            layer.ple_value      = create_tensor(tn(LLM_TENSOR_PLE_VALUE,      "weight", il), { ple_dim, n_embd }, 0);
             layer.ple_norm_key   = create_tensor(tn(LLM_TENSOR_PLE_NORM_KEY,   "weight", il), { hc_dim }, 0);
             layer.ple_norm_query = create_tensor(tn(LLM_TENSOR_PLE_NORM_QUERY, "weight", il), { hc_dim }, 0);
             layer.ple_norm_conv  = create_tensor(tn(LLM_TENSOR_PLE_NORM_CONV,  "weight", il), { hc_dim }, 0);
@@ -296,7 +345,6 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa
 
     ggml_tensor * inpL = build_inp_embd(model.tok_embd);
     cb(inpL, "model.input_embed", -1);
-    ggml_build_forward_expand(gf, inpL);
 
     auto * inp = build_inp_mem_hybrid();
 
@@ -313,13 +361,6 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa
     ggml_tensor * inp_pos     = build_inp_pos();
     ggml_tensor * inp_out_ids = build_inp_out_ids();
 
-    ggml_tensor * ple_emb = nullptr;
-    if (hparams.ple_n_heads > 0) {
-        ple_emb = build_inp_ple(mctx_hyb);
-        // make sure ple_emb and build_inp_embd are in the same graph split
-        ggml_build_forward_expand(gf, ple_emb);
-    }
-
     // the wide residual starts as hc identical copies of the embedding
     ggml_tensor * res_hc = ggml_repeat_4d(ctx0,
             ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens),
@@ -330,7 +371,7 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa
         res->t_layer_inp[il] = res_hc;
 
         if (hparams.is_ple(il)) {
-            res_hc = build_ple(inp->get_recr(), ple_emb, res_hc, il);
+            res_hc = build_ple(inp->get_recr(), mctx_hyb, res_hc, il);
         }
 
         ggml_tensor * inject = nullptr;
@@ -424,13 +465,17 @@ ggml_tensor * llama_model_qwen4exp::graph::build_norm_gated(
 // one mean-pooled indexer key scores each block; set_input resolves the cache layout
 class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i {
 public:
-    llm_graph_input_qsa(const llama_memory_hybrid_idx_context * mctx, uint32_t ratio, bool blk_bias) :
-        mctx(mctx), ratio(ratio), blk_bias(blk_bias) {}
+    llm_graph_input_qsa(
+            const llama_memory_hybrid_idx_context * mctx,
+            ggml_tensor * kq_mask,
+            uint32_t ratio,
+            uint32_t block_topk) :
+        mctx(mctx), kq_mask(kq_mask), ratio(ratio), block_topk(block_topk) {}
     virtual ~llm_graph_input_qsa() = default;
 
     void set_input(const llama_ubatch * ubatch) override {
-        mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch);
-        mctx->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias);
+        mctx->set_input_qsa(block_cells, block_pos, block_mask, selected,
+                kq_mask, ubatch, ratio, block_topk);
     }
 
     bool can_reuse(const llm_graph_params & params) override {
@@ -442,39 +487,33 @@ public:
         }
 
         const int64_t n_kv     = idx->get_n_kv();
-        const int64_t n_stream = mctx->get_n_stream();
-        const int64_t n_blocks = (n_kv + ratio - 1)/ratio;
+        const int64_t n_blocks = n_kv/ratio;
 
         bool res = true;
 
-        res &= params.ubatch.n_tokens % n_stream == 0;
-
-        res &= k_idxs->ne[0]    == params.ubatch.n_tokens;
-        res &= cell_blk->ne[0]  == n_kv;
-        res &= cell_blk->ne[1]  == n_stream;
-        res &= blk_cells->ne[0] == (int64_t) ratio*n_blocks;
-        res &= blk_pos->ne[0]   == 4*n_blocks*n_stream;
-        res &= bias->ne[0]      == (blk_bias ? n_blocks : n_kv);
-        res &= bias->ne[1]      == params.ubatch.n_tokens/n_stream;
+        res &= n_kv > (int64_t) block_topk*ratio + ratio - 1;
+        res &= block_cells->ne[1] == n_blocks;
+        res &= block_cells->ne[2] == params.ubatch.n_tokens;
+        res &= block_pos->ne[2]   == params.ubatch.n_tokens;
+        res &= block_mask->ne[1]  == params.ubatch.n_tokens;
+        res &= selected->ne[0]    == n_kv;
+        res &= selected->ne[1]    == params.ubatch.n_tokens;
 
         return res;
     }
 
-    // per stream: a cell index names a different token in each stream
-    ggml_tensor * k_idxs    = nullptr;   // I32 [n_tokens]
-    ggml_tensor * cell_blk  = nullptr;   // I32 [n_kv, n_stream]
-    ggml_tensor * blk_cells = nullptr;   // I32 [ratio*n_blocks, n_stream]
-    ggml_tensor * blk_pos   = nullptr;   // I32 [4*n_blocks*n_stream]
-    ggml_tensor * bias      = nullptr;   // F32 [n_blocks or n_kv, n_tokens/n_stream, n_stream]
+    ggml_tensor * block_cells = nullptr;
+    ggml_tensor * block_pos   = nullptr;
+    ggml_tensor * block_mask  = nullptr;
+    ggml_tensor * selected    = nullptr;
 
     const llama_memory_hybrid_idx_context * mctx;
+    ggml_tensor * kq_mask;
     const uint32_t ratio;
-
-    // the per-cell half of the bias is the attention mask, so only the per-block half is uploaded
-    const bool blk_bias;
+    const uint32_t block_topk;
 };
 
-ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
+ggml_tensor * llama_model_qwen4exp::graph::build_qsa_mask(
         const llama_memory_hybrid_idx_context * mctx_hyb,
         ggml_tensor *                           cur,
         ggml_tensor *                           inp_pos,
@@ -490,21 +529,13 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
 
     GGML_ASSERT(r > 0);
 
-    const int64_t n_blocks = (n_kv + r - 1)/r;
+    const int64_t n_blocks = n_kv/r;
+    const int64_t block_topk = hparams.indexer_top_k/r;
 
-    // build_attn_qsa and the KQ mask need the tokens to divide evenly across the streams
     const int64_t n_stream = mctx_hyb->get_n_stream();
     GGML_ASSERT(n_tokens % n_stream == 0);
     const int64_t n_tps = n_tokens/n_stream;
 
-    // only the "which block is visible" half of the bias varies per block
-    // the rest is the visible/not test the attention mask already carries, so upload the per-block half only: 1/ratio of the cells
-    // alibi writes distances instead of a mask and non-causal keeps future cells, so both opt out
-    // the mask also holds an mrope rule for the query's own position, but only 2d image positions can differ there
-    const bool blk_bias = kq_mask != nullptr &&
-        kq_mask->ne[0] == n_kv && kq_mask->ne[1] == n_tps && kq_mask->ne[3] == n_stream &&
-        cparams.causal_attn && !hparams.use_alibi;
-
     // nothing above depends on the layer, so the layers sharing a ratio share one input set
     llm_graph_input_qsa * inp = nullptr;
 
@@ -512,59 +543,29 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
     if (it != qsa_inps.end()) {
         inp = it->second;
     } else {
-        auto qsa = std::make_unique<llm_graph_input_qsa>(mctx_hyb, (uint32_t) r, blk_bias);
+        auto qsa = std::make_unique<llm_graph_input_qsa>(
+                mctx_hyb, kq_mask, (uint32_t) r, (uint32_t) block_topk);
 
-        qsa->k_idxs    = mctx_idx->build_input_k_idxs(ctx0, ubatch);
-        qsa->cell_blk  = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_stream);
-        qsa->blk_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, r*n_blocks, n_stream);
-        qsa->blk_pos   = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_blocks*n_stream);
-        qsa->bias      = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, blk_bias ? n_blocks : n_kv, n_tps, n_stream);
+        qsa->block_cells = ggml_new_tensor_3d(ctx0, GGML_TYPE_I32, r, n_blocks, n_tokens);
+        qsa->block_pos   = ggml_new_tensor_3d(ctx0, GGML_TYPE_I32, n_blocks, 4, n_tokens);
+        qsa->block_mask  = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_blocks, n_tokens);
+        qsa->selected    = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv, n_tokens);
 
-        ggml_set_input(qsa->cell_blk);
-        ggml_set_input(qsa->blk_cells);
-        ggml_set_input(qsa->blk_pos);
-        ggml_set_input(qsa->bias);
+        ggml_set_input(qsa->block_cells);
+        ggml_set_input(qsa->block_pos);
+        ggml_set_input(qsa->block_mask);
+        ggml_set_input(qsa->selected);
 
         inp = qsa.get();
         res->add_input(std::move(qsa));
         qsa_inps.emplace((uint32_t) r, inp);
     }
 
-    // cached indexer keys are raw: pooling precedes norm and rotation, so apply neither
-    ggml_tensor * k_raw = build_lora_mm(model.layers[il].index_k_proj, cur);
-    k_raw = ggml_reshape_3d(ctx0, k_raw, idx_dim, 1, n_tokens);
-    cb(k_raw, "indexer_k_raw", il);
-
-    ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, k_raw, inp->k_idxs, il));
+    kq_mask = inp->kq_mask;
 
     // one key head, so rows are contiguous. get_k gives [idx_dim, n_head_kv, n_kv, n_stream].
     ggml_tensor * k_all = mctx_idx->get_k(ctx0, il);
     k_all = ggml_view_3d(ctx0, k_all, idx_dim, n_kv, n_stream, k_all->nb[2], k_all->nb[3], 0);
-
-    // gathers per stream: blk_cells row s indexes stream s's own cells
-    ggml_tensor * members = ggml_get_rows(ctx0, k_all, inp->blk_cells);
-    members = ggml_reshape_4d(ctx0, members, idx_dim, r, n_blocks, n_stream);
-
-    // mean over the block members; r is small, so summing slices beats a transpose plus sum_rows
-    ggml_tensor * pooled = nullptr;
-    for (int64_t i = 0; i < r; ++i) {
-        ggml_tensor * slice = ggml_cont(ctx0,
-                ggml_view_3d(ctx0, members, idx_dim, n_blocks, n_stream,
-                        members->nb[2], members->nb[3], i*members->nb[1]));
-        pooled = pooled ? ggml_add(ctx0, pooled, slice) : slice;
-    }
-    pooled = ggml_scale(ctx0, pooled, 1.0f/(float) r);
-    cb(pooled, "indexer_k_pooled", il);
-
-    // rope wants [n_dims, n_head, n_tokens]: lay every stream's blocks flat, split after.
-    pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, 1, n_blocks*n_stream);
-    pooled = build_norm(pooled, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il);
-    pooled = ggml_rope_multi(ctx0, pooled, inp->blk_pos, nullptr,
-            n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,
-            ext_factor, attn_factor, beta_fast, beta_slow);
-    pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, n_blocks, n_stream);
-    cb(pooled, "indexer_k", il);
-
     ggml_tensor * q = build_lora_mm(model.layers[il].index_q_proj, cur);
     q = ggml_reshape_3d(ctx0, q, idx_dim, n_idx_h, n_tokens);
     q = build_norm(q, model.layers[il].index_q_norm, nullptr, LLM_NORM_RMS, il);
@@ -573,128 +574,86 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
             ext_factor, attn_factor, beta_fast, beta_slow);
     cb(q, "indexer_q", il);
 
-    // rectify each head dot product before the sum, as in the DeepSeek lightning indexer
-    // mul_mat matches ne[2], so the queries of stream s only meet the blocks of stream s
-    ggml_tensor * score = ggml_mul_mat(ctx0, pooled,
-            ggml_reshape_3d(ctx0, ggml_cont(ctx0, q), idx_dim, n_idx_h*n_tps, n_stream));
-    score = ggml_reshape_4d(ctx0, score, n_blocks, n_idx_h, n_tps, n_stream);
-    score = ggml_relu(ctx0, score);
-    score = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3));
-    score = ggml_sum_rows(ctx0, score);
-    score = ggml_reshape_3d(ctx0, score, n_blocks, n_tps, n_stream);
-    cb(score, "indexer_score", il);
+    const ggml_type activation_type = mctx_idx->type_k();
+    std::vector<ggml_tensor *> selected_streams;
+    selected_streams.reserve(n_stream);
 
-    // one value per block, so it is cheaper to bias here than after the cells are expanded
-    if (blk_bias) {
-        score = ggml_add(ctx0, score, inp->bias);
+    for (int64_t is = 0; is < n_stream; ++is) {
+        ggml_tensor * cache = ggml_view_2d(ctx0, k_all, idx_dim, n_kv,
+                k_all->nb[1], is*k_all->nb[2]);
+        ggml_tensor * block_cells = ggml_view_3d(ctx0, inp->block_cells, r, n_blocks, n_tps,
+                inp->block_cells->nb[1], inp->block_cells->nb[2], is*n_tps*inp->block_cells->nb[2]);
+        ggml_tensor * block_keys = ggml_get_rows(ctx0, cache,
+                ggml_reshape_1d(ctx0, block_cells, r*n_blocks*n_tps));
+        block_keys = ggml_reshape_4d(ctx0, block_keys, idx_dim, r, n_blocks, n_tps);
+        block_keys = ggml_cont(ctx0, ggml_transpose(ctx0, block_keys));
+        block_keys = ggml_mean(ctx0, block_keys);
+        block_keys = ggml_cont(ctx0, ggml_transpose(ctx0, block_keys));
+        block_keys = ggml_reshape_3d(ctx0, block_keys, idx_dim, 1, n_blocks*n_tps);
+
+        if (block_keys->type != activation_type) {
+            block_keys = ggml_cast(ctx0, block_keys, activation_type);
+            block_keys = ggml_cast(ctx0, block_keys, GGML_TYPE_F32);
+        }
+        block_keys = build_norm(block_keys, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il);
+
+        ggml_tensor * block_pos = ggml_view_3d(ctx0, inp->block_pos, n_blocks, 4, n_tps,
+                inp->block_pos->nb[1], inp->block_pos->nb[2], is*n_tps*inp->block_pos->nb[2]);
+        block_keys = ggml_rope_multi(ctx0, block_keys,
+                ggml_reshape_1d(ctx0, block_pos, n_blocks*4*n_tps), nullptr,
+                n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,
+                ext_factor, attn_factor, beta_fast, beta_slow);
+        cb(block_keys, "indexer_k", il);
+
+        block_keys = ggml_reshape_4d(ctx0, block_keys, idx_dim, n_blocks, 1, n_tps);
+        ggml_tensor * query = ggml_view_3d(ctx0, q, idx_dim, n_idx_h, n_tps,
+                q->nb[1], q->nb[2], is*n_tps*q->nb[2]);
+        query = ggml_reshape_4d(ctx0, query, idx_dim, n_idx_h, 1, n_tps);
+
+        ggml_tensor * scores = ggml_mul_mat(ctx0, block_keys, query);
+        ggml_mul_mat_set_prec(scores, GGML_PREC_F32);
+        scores = ggml_relu(ctx0, scores);
+        scores = ggml_sum_rows(ctx0, ggml_cont(ctx0, ggml_permute(ctx0, scores, 1, 0, 2, 3)));
+        scores = ggml_scale(ctx0, ggml_reshape_2d(ctx0, scores, n_blocks, n_tps),
+                1.0f/sqrtf((float) idx_dim));
+
+        ggml_tensor * block_mask = ggml_view_2d(ctx0, inp->block_mask, n_blocks, n_tps,
+                inp->block_mask->nb[1], is*n_tps*inp->block_mask->nb[1]);
+        scores = ggml_add(ctx0, scores, block_mask);
+        cb(scores, "indexer_score", il);
+
+        ggml_tensor * top_blocks = ggml_top_k(ctx0, scores, block_topk);
+        ggml_tensor * top_cells = ggml_get_rows(ctx0, block_cells, top_blocks);
+        top_cells = ggml_reshape_2d(ctx0, top_cells, r*block_topk, n_tps);
+
+        ggml_tensor * base_selected = ggml_view_2d(ctx0, inp->selected, n_kv, n_tps,
+                inp->selected->nb[1], is*n_tps*inp->selected->nb[1]);
+        base_selected = ggml_reshape_3d(ctx0, base_selected, 1, n_kv, n_tps);
+        ggml_tensor * selected_top = ggml_fill(ctx0, base_selected, 0.0f);
+        ggml_tensor * ones = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, r*block_topk, n_tps);
+        ones = ggml_fill(ctx0, ones, 1.0f);
+        selected_top = ggml_set_rows(ctx0, selected_top, ones, top_cells);
+        ggml_tensor * selected_stream = ggml_clamp(
+                ctx0, ggml_add(ctx0, base_selected, selected_top), 0.0f, 1.0f);
+        selected_streams.push_back(ggml_reshape_2d(ctx0, selected_stream, n_kv, n_tps));
     }
 
-    // every token of a block gets the block score; the budget is whole blocks, so top-k cuts on a block boundary
-    ggml_tensor * expanded = ggml_get_rows(ctx0,
-            ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), inp->cell_blk);
-    expanded = ggml_cont(ctx0, ggml_permute(ctx0, expanded, 1, 0, 2, 3));
-
-    if (blk_bias) {
-        // flash attention keeps the mask in f16; the scores are f32
-        ggml_tensor * mask = kq_mask->type == GGML_TYPE_F32 ? kq_mask : ggml_cast(ctx0, kq_mask, GGML_TYPE_F32);
-        expanded = ggml_add(ctx0, expanded, ggml_reshape_3d(ctx0, mask, n_kv, n_tps, n_stream));
-    } else {
-        expanded = ggml_add(ctx0, expanded, inp->bias);
+    ggml_tensor * selected = selected_streams[0];
+    for (int64_t is = 1; is < n_stream; ++is) {
+        selected = ggml_concat(ctx0, selected, selected_streams[is], 1);
     }
-    cb(expanded, "indexer_score_tokens", il);
+    selected = ggml_scale_bias(ctx0, selected, 1e30f, -1e30f);
 
-    // the reference returns indexer_top_k + compress_ratio - 1: whole blocks plus the tail
-    const int64_t width = std::min<int64_t>(n_kv, (int64_t) hparams.indexer_top_k + r - 1);
-
-    ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, width));
-
-    // build_attn_qsa reads [n_top_k, n_batch, 1, n_stream], matching the KQ mask.
-    top_k = ggml_reshape_4d(ctx0, top_k, width, n_tps, 1, n_stream);
-    cb(top_k, "indexer_top_k", il);
-
-    return top_k;
-}
-
-// Dense GQA self-attention restricted to the cells that top_k names.
-// The mask build below copies the MLA sparse path in llm_graph_context::build_attn.
-ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa(
-        llm_graph_input_attn_kv * inp,
-        ggml_tensor *             q_cur,
-        ggml_tensor *             k_cur,
-        ggml_tensor *             v_cur,
-        ggml_tensor *             top_k,
-        float                     kq_scale,
-        int                       il) {
-    // rotate q/k/v before they reach a quantized cache, as the dense path does. the indexer
-    // has already scored with its own query in build_qsa_top_k, so top_k is unaffected.
-    if (inp->self_k_rot) {
-        q_cur = llama_mul_mat_hadamard(ctx0, q_cur, inp->self_k_rot);
-        k_cur = llama_mul_mat_hadamard(ctx0, k_cur, inp->self_k_rot);
+    ggml_tensor * base_mask = ggml_reshape_2d(ctx0, kq_mask, n_kv, n_tokens);
+    if (base_mask->type != GGML_TYPE_F32) {
+        base_mask = ggml_cast(ctx0, base_mask, GGML_TYPE_F32);
     }
-
-    if (inp->self_v_rot) {
-        v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot);
+    ggml_tensor * mask = ggml_add(ctx0, base_mask, selected);
+    if (cparams.flash_attn) {
+        mask = ggml_cast(ctx0, mask, GGML_TYPE_F16);
     }
-
-    // these nodes are added to the graph together so that they are not reordered
-    // by doing so, the number of splits in the graph is reduced
-    // expand k later to enable rope fusion which directly writes into k-v cache
-    ggml_build_forward_expand(gf, q_cur);
-    ggml_build_forward_expand(gf, v_cur);
-    ggml_build_forward_expand(gf, k_cur);
-
-    const auto * mctx_cur = inp->mctx;
-
-    // store to KV cache
-    {
-        const auto & k_idxs = inp->get_k_idxs();
-        const auto & v_idxs = inp->get_v_idxs();
-
-        ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));
-        ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));
-    }
-
-    ggml_tensor * kq_mask = inp->get_kq_mask();
-
-    // prepare new kq mask - starts filled with -INFINITY
-    ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
-
-    // reshape KQ mask into tensor with rows of size 1:
-    // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream]
-    kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);
-
-    // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1]
-    ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0);
-
-    // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream]
-    // this will be our source of zero values for unmasking top k mask elements
-    ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]);
-    zeros = ggml_fill(ctx0, zeros, 0.0f);
-
-    // modify KQ mask by unmasking elements that are in top_k indices
-    // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1])
-    ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d);
-
-    // reshape to restore the original shape of KQ mask:
-    // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream]
-    kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0);
-
-    // combine with the original kq mask
-    kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask);
-
-    ggml_tensor * q = q_cur;
-    ggml_tensor * k = mctx_cur->get_k(ctx0, il);
-    ggml_tensor * v = mctx_cur->get_v(ctx0, il);
-
-    ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, kq_scale, il);
-    cb(cur, "kqv_out", il);
-
-    // the rotation is its own inverse, so undo it on the value side of the output
-    if (inp->self_v_rot) {
-        cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot);
-    }
-
-    return cur;
+    cb(mask, "qsa_mask", il);
+    return mask;
 }
 
 ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn(
@@ -707,10 +666,23 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn(
     const int64_t n_embd_head = hparams.n_embd_head_v();
     GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
 
-    // indexer reads the same block input as q/k/v; no cache or no ratio means dense
-    const bool qsa = mctx_hyb->get_idx() != nullptr && hparams.dsv4_compress_ratios[il] > 0;
+    const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx();
+    if (mctx_idx) {
+        ggml_tensor * index_k = build_lora_mm(model.layers[il].index_k_proj, cur);
+        index_k = ggml_reshape_3d(ctx0, index_k, hparams.indexer_head_size, 1, n_tokens);
+        cb(index_k, "indexer_k_raw", il);
+        ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, index_k, inp->get_k_idxs(), il));
+    }
 
-    ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, inp->get_kq_mask(), sections, il) : nullptr;
+    const int64_t ratio = hparams.dsv4_compress_ratios[il];
+    const bool qsa = mctx_idx && ratio > 0 &&
+            mctx_idx->get_n_kv() > (int64_t) hparams.indexer_top_k + ratio - 1;
+    if (qsa) {
+        inp->self_kq_mask_cnv = build_qsa_mask(
+                mctx_hyb, cur, inp_pos, inp->self_kq_mask, sections, il);
+    } else {
+        inp->self_kq_mask_cnv = inp->self_kq_mask;
+    }
 
     // Qwen3Next uses a single Q projection that outputs query + gate
     ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); // [ (n_embd_head * 2) * n_head, n_tokens ]
@@ -762,13 +734,9 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn(
 
     const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
 
-    if (top_k) {
-        cur = build_attn_qsa(inp, Qcur, Kcur, Vcur, top_k, kq_scale, il);
-    } else {
-        cur = build_attn(inp,
-                    nullptr, nullptr, nullptr,
-                    Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
-    }
+    cur = build_attn(inp,
+                nullptr, nullptr, nullptr,
+                Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
     cb(cur, "attn_pregate", il);
 
     ggml_tensor * gate_sigmoid = ggml_sigmoid(ctx0, gate);
@@ -1098,8 +1066,13 @@ ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at(
     return conv_input;
 }
 
-ggml_tensor * llama_model_qwen4exp::graph::build_inp_ple(
-        const llama_memory_hybrid_idx_context * mctx_hyb) {
+ggml_tensor * llama_model_qwen4exp::graph::build_ple(
+        llm_graph_input_rs * inp,
+        const llama_memory_hybrid_idx_context * mctx_hyb,
+        ggml_tensor *        hidden,
+        int                  il) {
+    const int64_t hc      = hparams.dsv4_hc_mult;
+    const int64_t hc_dim  = hc * n_embd;
     const int64_t n_heads = hparams.ple_n_heads;
 
     // the attention cells see every ubatch regardless of the layer types
@@ -1114,18 +1087,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_inp_ple(
     // gather then flatten the heads: get_rows lays the head dimension out slowest, as the reference does
     ggml_tensor * emb = ggml_get_rows(ctx0, model.per_layer_tok_embd, rows);
     emb = ggml_reshape_2d(ctx0, emb, hparams.ple_head_dim * n_heads, n_tokens);
-    cb(emb, "ple_embd", -1);
-
-    return emb;
-}
-
-ggml_tensor * llama_model_qwen4exp::graph::build_ple(
-        llm_graph_input_rs * inp,
-        ggml_tensor *        emb,
-        ggml_tensor *        hidden,
-        int                  il) {
-    const int64_t hc      = hparams.dsv4_hc_mult;
-    const int64_t hc_dim  = hc * n_embd;
+    cb(emb, "ple_embd", il);
 
     ggml_tensor * key   = build_lora_mm(model.layers[il].ple_key,   emb);
     ggml_tensor * value = build_lora_mm(model.layers[il].ple_value, emb);

tests/CMakeLists.txt

@@ -220,6 +220,16 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
         FIXTURES_REQUIRED generate-models
     )
 
+    llama_test(
+        test-recurrent-state-rollback
+        NAME test-recurrent-state-rollback-qwen4exp
+        LABEL main
+        ARGS -m "${MODEL_DIR}/qwen4exp-moe.gguf"
+    )
+    set_tests_properties(test-recurrent-state-rollback-qwen4exp PROPERTIES
+        FIXTURES_REQUIRED generate-models
+    )
+
     llama_test(
         test-recurrent-state-rollback
         NAME test-recurrent-state-rollback-nemotron-h

tests/test-alloc.cpp

@@ -19,8 +19,6 @@ struct dummy_backend_context {
     size_t alignment       = 8;
 
     ggml_backend_buffer_i              buffer_interface;
-    ggml_backend_device                device;
-    ggml_backend                       backend;
     std::vector<ggml_backend_buffer_t> buffers;
 
     size_t allocated_total() const {
@@ -85,27 +83,7 @@ static void dummy_backend_buffer_get_tensor(ggml_backend_buffer_t, const ggml_te
 
 static void dummy_backend_buffer_clear(ggml_backend_buffer_t, uint8_t) {}
 
-// ggml_backend_device interface
-
-static enum ggml_backend_dev_type dummy_backend_device_get_type(ggml_backend_dev_t) {
-    return GGML_BACKEND_DEVICE_TYPE_CPU;
-}
-
-static bool dummy_backend_device_supports_op(ggml_backend_dev_t, const ggml_tensor *) {
-    return true;
-}
-
-static bool dummy_backend_device_supports_buft(ggml_backend_dev_t device, ggml_backend_buffer_type_t buft) {
-    return device->context == buft->context;
-}
-
-// ggml_backend interface
-
-static const char * dummy_backend_get_name(ggml_backend_t) {
-    return "dummy_backend";
-}
-
-// dummy_backend
+// dummy_backend (not really a full backend, just provides what gallocr needs)
 
 struct dummy_backend {
     std::unique_ptr<dummy_backend_context> context;
@@ -126,16 +104,6 @@ static dummy_backend dummy_backend_init(size_t max_buffer_size, size_t alignment
     b.context->buffer_interface.get_tensor    = dummy_backend_buffer_get_tensor;
     b.context->buffer_interface.clear         = dummy_backend_buffer_clear;
 
-    b.context->device.context             = b.context.get();
-    b.context->device.iface.get_type      = dummy_backend_device_get_type;
-    b.context->device.iface.supports_op   = dummy_backend_device_supports_op;
-    b.context->device.iface.supports_buft = dummy_backend_device_supports_buft;
-
-    b.context->backend.context        = b.context.get();
-    b.context->backend.device         = &b.context->device;
-    b.context->backend.iface.get_name = dummy_backend_get_name;
-
-    b.buffer_type.device              = &b.context->device;
     b.buffer_type.context             = b.context.get();
     b.buffer_type.iface.get_name      = dummy_backend_buffer_type_get_name;
     b.buffer_type.iface.alloc_buffer  = dummy_backend_buffer_type_alloc_buffer;
@@ -615,41 +583,6 @@ static void test_reallocation() {
     }
 }
 
-static void test_backend_graph_optimize(ggml_backend_t, ggml_cgraph * graph, ggml_backend_graph_optimize_params * params) {
-    GGML_ASSERT(graph->n_nodes == 3);
-    params->add_alloc_dep(params->user_data, graph->nodes[0], graph->nodes[2]);
-}
-
-static bool graph_reuses_allocation(bool add_alloc_dep) {
-    auto [ctx, graph, ctx_ptr] = make_context();
-
-    ggml_tensor * x[4];
-    x[0] = make_input_with_size(ctx, 16);
-    x[1] = ggml_scale(ctx, x[0], 2.0f);
-    x[2] = ggml_scale(ctx, x[1], 2.0f);
-    x[3] = ggml_scale(ctx, x[2], 2.0f);
-
-    ggml_set_output(x[3]);
-    ggml_build_forward_expand(graph, x[3]);
-
-    dummy_backend backend = dummy_backend_init(SIZE_MAX);
-    if (add_alloc_dep) {
-        backend.context->backend.iface.graph_optimize = test_backend_graph_optimize;
-    }
-
-    ggml_backend_t             backend_ptr = &backend.context->backend;
-    ggml_backend_buffer_type_t buft        = &backend.buffer_type;
-    ggml_backend_sched_ptr     sched(ggml_backend_sched_new(&backend_ptr, &buft, 1, 8, false, true));
-    GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), graph));
-
-    return x[1]->data == x[2]->data;
-}
-
-static void test_graph_optimize_alloc_dep() {
-    GGML_ASSERT(graph_reuses_allocation(false));
-    GGML_ASSERT(!graph_reuses_allocation(true));
-}
-
 static void run(const char * name, void (*f)()) {
     printf("%s ", name);
     fflush(stdout);
@@ -671,6 +604,5 @@ int main() {
     run("test_multiple_buffer_types", test_multiple_buffer_types);
     run("test_buffer_size_zero", test_buffer_size_zero);
     run("test_reallocation", test_reallocation);
-    run("test_graph_optimize_alloc_dep", test_graph_optimize_alloc_dep);
     return 0;
 }

tests/test-backend-ops.cpp

@@ -8785,7 +8785,6 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         test_cases.emplace_back(new test_conv_transpose_2d({3, 2, 3, 1}, {2, 2, 1, 3}, 1, kernel_type));
         test_cases.emplace_back(new test_conv_transpose_2d({10, 10, 9, 1}, {3, 3, 1, 9}, 2, kernel_type));
         test_cases.emplace_back(new test_conv_transpose_2d({129, 63, 35, 1}, {3, 3, 48, 35}, 1, kernel_type));
-        test_cases.emplace_back(new test_conv_transpose_2d({10, 10, 9, 2}, {3, 3, 1, 9}, 2, kernel_type)); // for multiple batches
     }
 
     test_cases.emplace_back(new test_count_equal(GGML_TYPE_F32, {4,  500, 1, 1}));
@@ -9359,6 +9358,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     // test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 512, 262144, 9216, {1, 1}, {1, 1}));
 
     // test large experts*tokens
+    test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 512, 10, false, 64, 512, 64));
     for (bool b : {false, true}) {
         test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, b, 32, 1024, 16));
         test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 2, 2, b, 32, 8192, 64));
@@ -9367,12 +9367,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 16, b, 50, 200, 64));
     }
 
-    // For issue 27873
-    test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_IQ2_XXS, GGML_TYPE_F32, 1, 1, false, 1, 8192, 4096));
-
-    for (int k : {1, 63, 65}) {
-        test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 1, 1, false, 8, 16, k));
-    }
+    test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 1, 1, false, 8, 16, 1));
     test_cases.emplace_back(new test_mul_mat_id_fusion(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, false, 32, 32, 32, 3));
 
     // gpt-oss issue with Vulkan mmq_id

tests/test-llama-archs.cpp

@@ -79,7 +79,7 @@ static std::vector<llama_token> get_tokens(const uint32_t n_tokens, const uint32
     return ret;
 }
 
-static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
+static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe, const bool qwen_ple = false) {
     gguf_context_ptr ret(gguf_init_empty());
     llama_model_saver ms(arch, ret.get());
     const uint32_t n_ctx = 256;
@@ -252,8 +252,23 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
     if (arch == LLM_ARCH_QWEN4EXP) {
         ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT,    uint32_t(4));
         ms.add_kv(LLM_KV_HYPER_CONNECTION_LOW_RANK, uint32_t(8));
-        // without this the QSA layers fall back to dense and go uncovered
-        ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector<uint32_t>(n_layer, 4));
+        std::vector<uint32_t> ratios(n_layer, 0);
+        for (uint32_t il = 1; il < n_layer; il += 2) {
+            ratios[il] = 4;
+        }
+        ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, ratios);
+
+        if (qwen_ple) {
+            ms.add_kv(LLM_KV_PLE_LAYERS,                std::vector<uint32_t>({0}));
+            ms.add_kv(LLM_KV_PLE_NGRAM_SIZE,             uint32_t(2));
+            ms.add_kv(LLM_KV_PLE_HEADS_PER_NGRAM,        uint32_t(1));
+            ms.add_kv(LLM_KV_PLE_CONV_KERNEL,            uint32_t(2));
+            ms.add_kv(LLM_KV_PLE_EOS_TOKEN_ID,           uint32_t(1));
+            ms.add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, uint32_t(64));
+            ms.add_kv(LLM_KV_PLE_LAYER_MULTIPLIERS,      std::vector<uint64_t>({1, 3}));
+            ms.add_kv(LLM_KV_PLE_HEAD_OFFSETS,           std::vector<uint64_t>({0}));
+            ms.add_kv(LLM_KV_PLE_HEAD_VOCAB_SIZES,       std::vector<uint64_t>({16}));
+        }
     }
 
     // minimax-m3 keeps one indexer head per GQA head; the rest use a fixed 64 to match the fused
@@ -347,7 +362,8 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/)
 
 static std::pair<llama_model_ptr, llama_context_ptr> get_model_and_ctx(
         struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector<ggml_backend_dev_t> & devs,
-        const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false) {
+        const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false,
+        ggml_backend_sched_eval_callback cb_eval = nullptr, void * cb_eval_user_data = nullptr) {
     GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr));
     llama_model_params model_params = llama_model_default_params();
     model_params.progress_callback = silent_model_load_progress;
@@ -360,6 +376,8 @@ static std::pair<llama_model_ptr, llama_context_ptr> get_model_and_ctx(
     ctx_params.n_ctx = 0;
     ctx_params.n_threads = 4;
     ctx_params.n_threads_batch = 4;
+    ctx_params.cb_eval = cb_eval;
+    ctx_params.cb_eval_user_data = cb_eval_user_data;
     if (!encode) {
         ctx_params.n_ubatch = 64;
     }
@@ -412,6 +430,54 @@ static std::vector<float> get_logits(
     return ret;
 }
 
+struct qwen4_qsa_mask_check {
+    int64_t n_seen = 0;
+    int64_t n_tokens_seen = 0;
+    bool ok = true;
+};
+
+static bool check_qwen4_qsa_mask(ggml_tensor * tensor, bool ask, void * user_data) {
+    if (strncmp(tensor->name, "qsa_mask", 8) != 0) {
+        return false;
+    }
+    if (ask) {
+        return true;
+    }
+
+    auto * check = (qwen4_qsa_mask_check *) user_data;
+    const int64_t n_kv = tensor->ne[0];
+    const int64_t n_tokens = tensor->ne[1];
+    std::vector<float> mask(ggml_nelements(tensor));
+    if (tensor->type == GGML_TYPE_F32) {
+        ggml_backend_tensor_get(tensor, mask.data(), 0, ggml_nbytes(tensor));
+    } else {
+        GGML_ASSERT(tensor->type == GGML_TYPE_F16);
+        std::vector<ggml_fp16_t> mask_f16(ggml_nelements(tensor));
+        ggml_backend_tensor_get(tensor, mask_f16.data(), 0, ggml_nbytes(tensor));
+        for (size_t i = 0; i < mask.size(); ++i) {
+            mask[i] = ggml_fp16_to_fp32(mask_f16[i]);
+        }
+    }
+
+    for (int64_t it = 0; it < n_tokens; ++it) {
+        const int64_t n_visible = check->n_tokens_seen + it + 1;
+        const int64_t n_complete = n_visible/4;
+        const int64_t expected = n_complete <= 2 ? n_visible : 8 + n_visible%4;
+        int64_t actual = 0;
+        for (int64_t ikv = 0; ikv < n_kv; ++ikv) {
+            actual += mask[it*n_kv + ikv] > -1e20f;
+        }
+        if (actual != expected) {
+            fprintf(stderr, "Qwen4 QSA mask row %lld selects %lld tokens, expected %lld\n",
+                    (long long) (check->n_tokens_seen + it), (long long) actual, (long long) expected);
+        }
+        check->ok = check->ok && actual == expected;
+    }
+    check->n_tokens_seen += n_tokens;
+    check->n_seen++;
+    return true;
+}
+
 static bool moe_mandatory(const llm_arch arch) {
     switch (arch) {
         case LLM_ARCH_LLAMA4:
@@ -686,6 +752,32 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg
             if (arch == LLM_ARCH_BAILINGMOE3) {
                 GGML_ASSERT(gguf_remove_key(gguf_ctx.get(), "bailingmoe3.kda.safe_gate") >= 0);
             }
+            if (arch == LLM_ARCH_QWEN4EXP) {
+                qwen4_qsa_mask_check check;
+                auto model_and_ctx = get_model_and_ctx(
+                        gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, false,
+                        check_qwen4_qsa_mask, &check);
+                get_logits(model_and_ctx.first.get(), model_and_ctx.second.get(), tokens);
+                GGML_ASSERT(check.ok && check.n_seen > 0);
+
+                gguf_context_ptr gguf_ctx_ple = get_gguf_ctx(arch, moe, true);
+                auto model_and_ctx_ple = get_model_and_ctx(gguf_ctx_ple.get(), nullptr, seed, {});
+                const std::vector<float> logits_ple = get_logits(
+                        model_and_ctx_ple.first.get(), model_and_ctx_ple.second.get(), tokens);
+
+                FILE * file_ple = tmpfile();
+                GGML_ASSERT(file_ple);
+                llama_model_saver saver_ple(model_and_ctx_ple.first.get());
+                saver_ple.add_kv_from_model();
+                saver_ple.add_tensors_from_model();
+                saver_ple.save(file_ple);
+                rewind(file_ple);
+
+                auto model_and_ctx_ple_saved = get_model_and_ctx(nullptr, file_ple, seed, {});
+                const std::vector<float> logits_ple_saved = get_logits(
+                        model_and_ctx_ple_saved.first.get(), model_and_ctx_ple_saved.second.get(), tokens);
+                GGML_ASSERT(logits_ple == logits_ple_saved);
+            }
             std::pair<llama_model_ptr, llama_context_ptr> model_and_ctx_cpu;
             std::vector<float> logits_cpu;
             for (device_config & dc : dev_configs) {

tools/llama-bench/README.md

@@ -67,7 +67,6 @@ test parameters:
   -nkvo, --no-kv-offload <0|1>              (default: 0)
   -fa, --flash-attn <on|off|auto>           (default: auto)
   -dev, --device <dev0/dev1/...>            (default: auto)
-  --tensor-read-lazy <on|auto|off>          (default: auto)
   -mmp, --mmap <0|1>                        (DEPRECATED IN FAVOUR OF --load-mode)
   -dio, --direct-io <0|1>                   (DEPRECATED IN FAVOUR OF --load-mode)
   -embd, --embeddings <0|1>                 (default: 0)

tools/llama-bench/llama-bench.cpp

@@ -271,19 +271,6 @@ static const char * split_mode_str(llama_split_mode mode) {
     }
 }
 
-static const char * lazy_mode_str(llama_lazy_mode mode) {
-    switch (mode) {
-        case LLAMA_LAZY_MODE_OFF:
-            return "off";
-        case LLAMA_LAZY_MODE_AUTO:
-            return "auto";
-        case LLAMA_LAZY_MODE_ON:
-            return "on";
-        default:
-            GGML_ABORT("invalid tensor read lazy mode");
-    }
-}
-
 static std::string pair_str(const std::pair<int, int> & p) {
     static char buf[32];
     snprintf(buf, sizeof(buf), "%d,%d", p.first, p.second);
@@ -354,7 +341,6 @@ struct cmd_params {
     std::vector<int>                 n_cpu_moe;
     std::vector<llama_split_mode>    split_mode;
     std::vector<llama_load_mode>     load_mode;
-    std::vector<llama_lazy_mode>     lazy_mode;
     std::vector<int>                 main_gpu;
     std::vector<bool>                no_kv_offload;
     std::vector<llama_flash_attn_type> flash_attn;
@@ -399,7 +385,6 @@ static const cmd_params cmd_params_defaults = {
     /* n_cpu_moe            */ { 0 },
     /* split_mode           */ { LLAMA_SPLIT_MODE_LAYER },
     /* load_mode            */ { LLAMA_LOAD_MODE_AUTO },
-    /* lazy_mode            */ { LLAMA_LAZY_MODE_AUTO },
     /* main_gpu             */ { 0 },
     /* no_kv_offload        */ { false },
     /* flash_attn           */ { LLAMA_FLASH_ATTN_TYPE_AUTO },
@@ -475,7 +460,6 @@ static void print_usage(int /* argc */, char ** argv) {
     printf("  -fa, --flash-attn <on|off|auto>                   (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str());
     printf("  -dev, --device <dev0/dev1/...>                    (default: auto)\n");
     printf("  -lm, --load-mode <auto|none|mmap|mlock|mmap+mlock|dio> (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str());
-    printf("  --tensor-read-lazy <on|auto|off>                  (default: %s)\n", join(transform_to_str(cmd_params_defaults.lazy_mode, lazy_mode_str), ",").c_str());
     printf("  -mmp, --mmap <0|1>                                (DEPRECATED IN FAVOUR OF --load-mode)\n");
     printf("  -dio, --direct-io <0|1>                           (DEPRECATED IN FAVOUR OF --load-mode)\n");
     printf("  -embd, --embeddings <0|1>                         (default: %s)\n", join(cmd_params_defaults.embeddings, ",").c_str());
@@ -802,32 +786,6 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
                     break;
                 }
                 params.load_mode.insert(params.load_mode.end(), modes.begin(), modes.end());
-            } else if (arg == "--tensor-read-lazy") {
-                if (++i >= argc) {
-                    invalid_param = true;
-                    break;
-                }
-                auto p = string_split<std::string>(argv[i], split_delim);
-
-                std::vector<llama_lazy_mode> modes;
-                for (const auto & m : p) {
-                    llama_lazy_mode mode;
-                    if (m == "on") {
-                        mode = LLAMA_LAZY_MODE_ON;
-                    } else if (m == "auto") {
-                        mode = LLAMA_LAZY_MODE_AUTO;
-                    } else if (m == "off") {
-                        mode = LLAMA_LAZY_MODE_OFF;
-                    } else {
-                        invalid_param = true;
-                        break;
-                    }
-                    modes.push_back(mode);
-                }
-                if (invalid_param) {
-                    break;
-                }
-                params.lazy_mode.insert(params.lazy_mode.end(), modes.begin(), modes.end());
             } else if (arg == "-mg" || arg == "--main-gpu") {
                 if (++i >= argc) {
                     invalid_param = true;
@@ -1179,9 +1137,6 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
     if (params.load_mode.empty()) {
         params.load_mode = cmd_params_defaults.load_mode;
     }
-    if (params.lazy_mode.empty()) {
-        params.lazy_mode = cmd_params_defaults.lazy_mode;
-    }
     if (params.main_gpu.empty()) {
         params.main_gpu = cmd_params_defaults.main_gpu;
     }
@@ -1248,7 +1203,6 @@ struct cmd_params_instance {
     int                n_cpu_moe;
     llama_split_mode   split_mode;
     llama_load_mode    load_mode;
-    llama_lazy_mode    lazy_mode;
     int                main_gpu;
     bool               no_kv_offload;
     llama_flash_attn_type flash_attn;
@@ -1270,7 +1224,6 @@ struct cmd_params_instance {
         }
         mparams.split_mode    = split_mode;
         mparams.load_mode     = load_mode;
-        mparams.lazy_mode     = lazy_mode;
         mparams.main_gpu      = main_gpu;
         mparams.tensor_split  = tensor_split.data();
         mparams.no_host       = no_host;
@@ -1318,8 +1271,7 @@ struct cmd_params_instance {
         return model == other.model && n_gpu_layers == other.n_gpu_layers && n_cpu_moe == other.n_cpu_moe &&
                split_mode == other.split_mode &&
                main_gpu == other.main_gpu && tensor_split == other.tensor_split &&
-               load_mode == other.load_mode && lazy_mode == other.lazy_mode &&
-               devices == other.devices && no_host == other.no_host &&
+               load_mode == other.load_mode && devices == other.devices && no_host == other.no_host &&
                vec_tensor_buft_override_equal(tensor_buft_overrides, other.tensor_buft_overrides);
     }
 
@@ -1353,7 +1305,6 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
     for (const auto & ncmoe : params.n_cpu_moe)
     for (const auto & sm : params.split_mode)
     for (const auto & lm : params.load_mode)
-    for (const auto & lzm : params.lazy_mode)
     for (const auto & mg : params.main_gpu)
     for (const auto & devs : params.devices)
     for (const auto & ts : params.tensor_split)
@@ -1393,7 +1344,6 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
                 /* .n_cpu_moe             = */ ncmoe,
                 /* .split_mode            = */ sm,
                 /* .load_mode             = */ lm,
-                /* .lazy_mode             = */ lzm,
                 /* .main_gpu              = */ mg,
                 /* .no_kv_offload         = */ nkvo,
                 /* .flash_attn            = */ fa,
@@ -1430,7 +1380,6 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
                 /* .n_cpu_moe             = */ ncmoe,
                 /* .split_mode            = */ sm,
                 /* .load_mode             = */ lm,
-                /* .lazy_mode             = */ lzm,
                 /* .main_gpu              = */ mg,
                 /* .no_kv_offload         = */ nkvo,
                 /* .flash_attn            = */ fa,
@@ -1467,7 +1416,6 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
                 /* .n_cpu_moe             = */ ncmoe,
                 /* .split_mode            = */ sm,
                 /* .load_mode             = */ lm,
-                /* .lazy_mode             = */ lzm,
                 /* .main_gpu              = */ mg,
                 /* .no_kv_offload         = */ nkvo,
                 /* .flash_attn            = */ fa,
@@ -1509,7 +1457,6 @@ struct test {
     int                      n_cpu_moe;
     llama_split_mode         split_mode;
     llama_load_mode          load_mode;
-    llama_lazy_mode          lazy_mode;
     int                      main_gpu;
     bool                     no_kv_offload;
     llama_flash_attn_type    flash_attn;
@@ -1549,7 +1496,6 @@ struct test {
         n_cpu_moe      = inst.n_cpu_moe;
         split_mode     = inst.split_mode;
         load_mode      = inst.load_mode;
-        lazy_mode      = inst.lazy_mode;
         main_gpu       = inst.main_gpu;
         no_kv_offload  = inst.no_kv_offload;
         flash_attn     = inst.flash_attn;
@@ -1617,8 +1563,7 @@ struct test {
             "n_ubatch",       "n_threads",      "cpu_mask",      "cpu_strict",     "poll",
             "type_k",         "type_v",         "n_gpu_layers",  "n_cpu_moe",      "split_mode",
             "main_gpu",       "no_kv_offload",  "flash_attn",    "devices",        "tensor_split",
-            "tensor_buft_overrides",            "load_mode",     "lazy_mode",
-            "embeddings",
+            "tensor_buft_overrides",            "load_mode",     "embeddings",
             "no_op_offload",  "no_host",        "fit_target",    "fit_min_ctx",
             "n_prompt",       "n_gen",          "n_depth",
             "test_time",      "avg_ns",         "stddev_ns",     "avg_ts",         "stddev_ts"
@@ -1643,7 +1588,7 @@ struct test {
         if (field == "avg_ts" || field == "stddev_ts") {
             return FLOAT;
         }
-        if (field == "load_mode" || field == "lazy_mode") {
+        if (field == "load_mode") {
             return STRING;
         }
         return STRING;
@@ -1713,7 +1658,6 @@ struct test {
                                             tensor_split_str,
                                             tensor_buft_overrides_str,
                                             llama_load_mode_name(load_mode),
-                                            lazy_mode_str(lazy_mode),
                                             std::to_string(embeddings),
                                             std::to_string(no_op_offload),
                                             std::to_string(no_host),
@@ -2028,9 +1972,6 @@ struct markdown_printer : public printer {
         if (params.load_mode.size() > 1 || params.load_mode != cmd_params_defaults.load_mode) {
             fields.emplace_back("load_mode");
         }
-        if (params.lazy_mode.size() > 1 || params.lazy_mode != cmd_params_defaults.lazy_mode) {
-            fields.emplace_back("lazy_mode");
-        }
         if (params.embeddings.size() > 1 || params.embeddings != cmd_params_defaults.embeddings) {
             fields.emplace_back("embeddings");
         }

vendor/cpp-httplib/httplib.h

@@ -8,8 +8,8 @@
 #ifndef CPPHTTPLIB_HTTPLIB_H
 #define CPPHTTPLIB_HTTPLIB_H
 
-#define CPPHTTPLIB_VERSION "0.54.1"
-#define CPPHTTPLIB_VERSION_NUM "0x003601"
+#define CPPHTTPLIB_VERSION "0.53.1"
+#define CPPHTTPLIB_VERSION_NUM "0x003501"
 
 #ifdef _WIN32
 #if defined(_WIN32_WINNT) && _WIN32_WINNT < 0x0A00
@@ -134,16 +134,6 @@
 #define CPPHTTPLIB_FORM_URL_ENCODED_PAYLOAD_MAX_LENGTH 8192
 #endif
 
-#ifndef CPPHTTPLIB_STATIC_FILE_COMPRESSION_MIN_LENGTH
-// 1400 rather than a round number: a body that already fits in one 1500-byte
-// MTU gains nothing from being made smaller.
-#define CPPHTTPLIB_STATIC_FILE_COMPRESSION_MIN_LENGTH 1400
-#endif
-
-#ifndef CPPHTTPLIB_STATIC_FILE_COMPRESSION_MAX_LENGTH
-#define CPPHTTPLIB_STATIC_FILE_COMPRESSION_MAX_LENGTH (4 * 1024 * 1024) // 4MB
-#endif
-
 #ifndef CPPHTTPLIB_RANGE_MAX_COUNT
 #define CPPHTTPLIB_RANGE_MAX_COUNT 1024
 #endif
@@ -1439,16 +1429,9 @@ public:
   DataSink &operator=(DataSink &&) = delete;
 
   std::function<bool(const char *data, size_t data_len)> write;
-
-  // Only `write` is mandatory. The rest are defaulted so that a provider
-  // calling one on a writer that does not set it gets sensible behaviour
-  // rather than std::bad_function_call thrown from a worker thread. Capturing
-  // `this` is safe: DataSink is neither copyable nor movable.
-  std::function<bool()> is_writable = []() { return true; };
-  std::function<void()> done = []() {};
-  std::function<void(const Headers &trailer)> done_with_trailer =
-      [this](const Headers & /*trailer*/) { done(); };
-
+  std::function<bool()> is_writable;
+  std::function<void()> done;
+  std::function<void(const Headers &trailer)> done_with_trailer;
   std::ostream os;
 
 private:
@@ -1533,10 +1516,7 @@ make_file_body(const std::string &filepath) {
       auto to_read = (std::min)(sizeof(buf), length);
       f.read(buf, static_cast<std::streamsize>(to_read));
       auto n = static_cast<size_t>(f.gcount());
-      // The file is shorter than the size make_file_body() measured, which the
-      // caller has already committed to as Content-Length. The body cannot be
-      // completed, so fail as every other error here does.
-      if (n == 0) { return false; }
+      if (n == 0) { break; }
       if (!sink.write(buf, n)) { return false; }
       length -= n;
     }
@@ -1743,14 +1723,6 @@ struct Request {
 #endif
 };
 
-namespace detail {
-
-// Declared up here, away from the rest of the compression helpers, because
-// `Response` stores one.
-enum class EncodingType { None = 0, Gzip, Brotli, Zstd };
-
-} // namespace detail
-
 struct Response {
   std::string version;
   int status = -1;
@@ -1816,11 +1788,6 @@ struct Response {
   bool content_provider_success_ = false;
   std::string file_content_path_;
   std::string file_content_content_type_;
-
-  // Content coding chosen for a file-backed content provider, decided once
-  // where the file is opened so that the ETag and the body cannot disagree.
-  // `EncodingType::None` for every other kind of response.
-  detail::EncodingType file_content_encoding_ = detail::EncodingType::None;
 };
 
 enum class Error {
@@ -1860,7 +1827,6 @@ enum class Error {
   InvalidRangeHeader,
   UnsupportedContentEncoding,
   WebSocketHandshake,
-  UserCallbackException,
 
   // For internal use only
   SSLPeerCouldBeClosed_,
@@ -2054,10 +2020,6 @@ private:
 
 int close_socket(socket_t sock) noexcept;
 
-bool is_accept_resource_error();
-
-bool is_accept_transient_error();
-
 ssize_t write_headers(Stream &strm, const Headers &headers);
 
 bool set_socket_opt_time(socket_t sock, int level, int optname, time_t sec,
@@ -2145,17 +2107,6 @@ public:
   Server &Delete(const std::string &pattern, HandlerWithContentReader handler);
   Server &Options(const std::string &pattern, Handler handler);
 
-  // Register a handler for an HTTP method outside the built-in set (e.g. the
-  // WebDAV methods from RFC 4918). Registering a method here is what makes the
-  // server accept it; an unregistered method is still rejected with 400.
-  // `method` must be a valid HTTP method token and must not be one of the
-  // built-in methods, which have their own registration functions above. A
-  // rejected registration makes is_valid() return false, so listen() fails.
-  Server &CustomRoute(const std::string &method, const std::string &pattern,
-                      Handler handler);
-  Server &CustomRoute(const std::string &method, const std::string &pattern,
-                      HandlerWithContentReader handler);
-
   Server &WebSocket(const std::string &pattern, WebSocketHandler handler);
   Server &WebSocket(const std::string &pattern, WebSocketHandler handler,
                     SubProtocolSelector sub_protocol_selector);
@@ -2223,10 +2174,6 @@ public:
 
   Server &set_payload_max_length(size_t length);
 
-  Server &set_static_file_compression(bool on);
-  Server &set_static_file_compression_min_length(size_t length);
-  Server &set_static_file_compression_max_length(size_t length);
-
   Server &set_websocket_ping_interval(time_t sec);
   template <class Rep, class Period>
   Server &set_websocket_ping_interval(
@@ -2255,35 +2202,6 @@ protected:
                        const std::function<void(Request &)> &setup_request,
                        bool *websocket_upgraded = nullptr);
 
-  // Runs the per-connection serving loop and stops an exception thrown by a
-  // user callback from escaping the worker thread.
-  //
-  // process_request() wraps only routing() in a try/catch. Content providers,
-  // the post-routing, error, logging and expect-100 handlers and WebSocket
-  // handlers all run outside it, and the task queue calls the job without a
-  // catch, so an exception from any of those would terminate the process.
-  //
-  // No 500 is possible here: by the time a content provider runs, the status
-  // line and headers are already on the wire. Report it through the error
-  // logger and drop the connection, which is what the peer observes either
-  // way. Other connections are unaffected.
-  template <typename Serve> bool serve_guarded(Serve &&serve) const {
-#ifdef CPPHTTPLIB_NO_EXCEPTIONS
-    return serve();
-#else
-    try {
-      return serve();
-    } catch (...) {
-      // The error logger is a user callback too, so it must not be able to
-      // throw the guard back open.
-      try {
-        output_error_log(Error::UserCallbackException, nullptr);
-      } catch (...) {}
-      return false;
-    }
-#endif
-  }
-
   std::atomic<socket_t> svr_sock_{INVALID_SOCKET};
 
   std::vector<std::string> trusted_proxies_;
@@ -2297,11 +2215,6 @@ protected:
   time_t idle_interval_sec_ = CPPHTTPLIB_IDLE_INTERVAL_SECOND;
   time_t idle_interval_usec_ = CPPHTTPLIB_IDLE_INTERVAL_USECOND;
   size_t payload_max_length_ = CPPHTTPLIB_PAYLOAD_MAX_LENGTH;
-  bool static_file_compression_ = false;
-  size_t static_file_compression_min_length_ =
-      CPPHTTPLIB_STATIC_FILE_COMPRESSION_MIN_LENGTH;
-  size_t static_file_compression_max_length_ =
-      CPPHTTPLIB_STATIC_FILE_COMPRESSION_MAX_LENGTH;
   time_t websocket_ping_interval_sec_ =
       CPPHTTPLIB_WEBSOCKET_PING_INTERVAL_SECOND;
   int websocket_max_missed_pongs_ = CPPHTTPLIB_WEBSOCKET_MAX_MISSED_PONGS;
@@ -2313,21 +2226,9 @@ private:
       std::vector<std::pair<std::unique_ptr<detail::MatcherBase>,
                             HandlerWithContentReader>>;
 
-  // Both handler tables for one custom method live in a single entry, so that
-  // routing() needs only one map lookup per request to reach either of them.
-  struct CustomHandlerEntry {
-    Handlers handlers;
-    HandlersForContentReader handlers_for_content_reader;
-  };
-  using CustomHandlers = std::map<std::string, CustomHandlerEntry>;
-
   static std::unique_ptr<detail::MatcherBase>
   make_matcher(const std::string &pattern);
 
-  static const std::set<std::string> &builtin_methods();
-  CustomHandlerEntry *custom_entry_for_registration(const std::string &method);
-  const CustomHandlerEntry *find_custom_entry(const std::string &method) const;
-
   template <typename H>
   Server &add_handler(
       std::vector<std::pair<std::unique_ptr<detail::MatcherBase>, H>> &handlers,
@@ -2358,10 +2259,6 @@ private:
       const HandlersForContentReader &handlers) const;
 
   bool parse_request_line(const char *s, Request &req) const;
-  detail::EncodingType static_file_encoding(const Request &req,
-                                            const std::string &content_type,
-                                            size_t length) const;
-  bool apply_static_file_compression(const Request &req, Response &res) const;
   void apply_ranges(const Request &req, Response &res,
                     std::string &content_type, std::string &boundary) const;
   bool write_response(Stream &strm, bool close_connection, Request &req,
@@ -2395,10 +2292,6 @@ private:
   std::atomic<bool> is_running_{false};
   std::atomic<bool> is_decommissioned{false};
 
-  // Set when CustomRoute() refuses a registration. Written before listen(),
-  // read by is_valid() on the same thread, so it needs no synchronization.
-  bool has_invalid_registration_ = false;
-
   struct MountPointEntry {
     std::string mount_point;
     std::string base_dir;
@@ -2420,7 +2313,6 @@ private:
   Handlers delete_handlers_;
   HandlersForContentReader delete_handlers_for_content_reader_;
   Handlers options_handlers_;
-  CustomHandlers custom_handlers_;
 
   struct WebSocketHandlerEntry {
     std::unique_ptr<detail::MatcherBase> matcher;
@@ -3608,16 +3500,6 @@ void split(const char *b, const char *e, char d,
 void split(const char *b, const char *e, char d, size_t m,
            std::function<void(const char *, const char *)> fn);
 
-bool split_find(const char *b, const char *e, char d,
-                std::function<bool(const char *, const char *)> fn);
-
-bool has_header_token(const Headers &headers, const std::string &key,
-                      const std::string &token);
-
-std::string websocket_accept_key(const std::string &client_key);
-
-bool is_websocket_upgrade(const Request &req);
-
 bool process_client_socket(
     socket_t sock, time_t read_timeout_sec, time_t read_timeout_usec,
     time_t write_timeout_sec, time_t write_timeout_usec,
@@ -3638,9 +3520,6 @@ socket_t create_client_socket(const std::string &host, const std::string &ip,
 const char *get_header_value(const Headers &headers, const std::string &key,
                              const char *def, size_t id);
 
-std::string get_combined_header_value(const Headers &headers,
-                                      const std::string &key);
-
 std::string params_to_query_str(const Params &params);
 
 void parse_query_text(const char *data, std::size_t size, Params &params);
@@ -3655,13 +3534,11 @@ bool parse_range_header(const std::string &s, Ranges &ranges);
 bool parse_accept_header(const std::string &s,
                          std::vector<std::string> &content_types);
 
-void parse_disposition_params(const std::string &s, Params &params);
-
 ssize_t send_socket(socket_t sock, const void *ptr, size_t size, int flags);
 
 ssize_t read_socket(socket_t sock, void *ptr, size_t size, int flags);
 
-EncodingType encoding_type(const Request &req, const std::string &content_type);
+enum class EncodingType { None = 0, Gzip, Brotli, Zstd };
 
 EncodingType encoding_type(const Request &req, const Response &res);
 
@@ -4441,11 +4318,6 @@ private:
   int unacked_pings_ = 0;
   std::atomic<bool> closed_{false};
   std::mutex write_mutex_;
-  // Owned by whichever thread is parsing frames off strm_. Only one thread
-  // may do so: read_websocket_frame() reads a payload until it has the whole
-  // declared length, so a second parser stealing bytes silently corrupts the
-  // message the first one is assembling.
-  std::mutex read_mutex_;
   std::thread ping_thread_;
   std::mutex ping_mutex_;
   std::condition_variable ping_cv_;