metal : add CONV_2D_DW (depthwise convolution) support (#21565)

* preliminary extract common header

* op source split

* split metallib into 8 libs && load in parallel

* derive kernel->library routing from functionNames

* x-macro lib list + underscore filenames, dedup QK_NL, MRC fixes

* op source split 8 to 20

* improve robustness of source fallback

* clean up

* change bool -> atomic_bool

* only prepend headers that source actually includes

* no semaphore, use GCD global queue

* dedup library compile path, fix NSError lifetime, rename gla

* relocate upstream concat/rope_back/repeat kernel changes into split files

* move ggml-common.h from common.h into dequantize.h to shrink binary size

---------

Co-authored-by: lvyichen <lvyichen@stepfun.com>

.github/workflows/ui-publish.yml

@@ -73,3 +73,4 @@ jobs:
           hf buckets rm ggml-org/${{ env.HF_BUCKET_NAME }}/index.html --yes 2>/dev/null || true
           hf buckets rm ggml-org/${{ env.HF_BUCKET_NAME }}/bundle.js --yes 2>/dev/null || true
           hf buckets rm ggml-org/${{ env.HF_BUCKET_NAME }}/bundle.css --yes 2>/dev/null || true
+          hf buckets rm ggml-org/${{ env.HF_BUCKET_NAME }}/loading.html --yes 2>/dev/null || true

common/arg.cpp

@@ -488,15 +488,12 @@ void common_models_handler_apply(common_models_handler & handler, common_params
         task.opts       = opts;
         tasks.push_back(task);
     }
-
-    bool had_spec_url = false;
     if (!params.speculative.draft.mparams.url.empty()) {
         common_download_task task;
         task.url        = params.speculative.draft.mparams.url;
         task.local_path = params.speculative.draft.mparams.path;
         task.opts       = opts;
         tasks.push_back(task);
-        had_spec_url = true;
     }
 
     // handle hf_plan tasks
@@ -516,18 +513,6 @@ void common_models_handler_apply(common_models_handler & handler, common_params
             });
         }
     };
-
-    // handle plan_spec (e.g. --spec-draft-hf)
-    if (!plan_spec.model_files.empty() && !had_spec_url) {
-        add_tasks(plan_spec.model_files, plan_spec.primary, params.speculative.draft.mparams);
-        had_spec_url = true;
-    }
-
-    // handle vocoder plan (e.g. --hf-repo-v)
-    if (!plan_voc.model_files.empty()) {
-        add_tasks(plan_voc.model_files, plan_voc.primary, params.vocoder.model);
-    }
-
     if (!plan.model_files.empty()) {
         add_tasks(plan.model_files, plan.primary, params.model);
     }
@@ -536,7 +521,7 @@ void common_models_handler_apply(common_models_handler & handler, common_params
             params.mmproj.path = hf_cache::finalize_file(plan.mmproj);
         });
     }
-    if (!plan.mtp.local_path.empty() && !had_spec_url) {
+    if (!plan.mtp.local_path.empty()) {
         tasks.emplace_back(plan.mtp, opts, [&]() {
             // only fall back to the discovered MTP head when no draft was explicitly provided
             if (params.speculative.draft.mparams.empty()) {
@@ -555,6 +540,16 @@ void common_models_handler_apply(common_models_handler & handler, common_params
         });
     }
 
+    // handle plan_spec (e.g. --spec-draft-hf)
+    if (!plan_spec.model_files.empty()) {
+        add_tasks(plan_spec.model_files, plan_spec.primary, params.speculative.draft.mparams);
+    }
+
+    // handle vocoder plan (e.g. --hf-repo-v)
+    if (!plan_voc.model_files.empty()) {
+        add_tasks(plan_voc.model_files, plan_voc.primary, params.vocoder.model);
+    }
+
     // run all tasks in parallel
     if (!params.offline) {
         // if duplicated files are found, only download once (but still call on_done for each task)
@@ -567,7 +562,6 @@ void common_models_handler_apply(common_models_handler & handler, common_params
         }
         std::vector<common_download_task> unique_tasks_vec;
         for (auto & pair : unique_tasks) {
-            LOG_DBG("download task: %s -> %s\n", pair.second->url.c_str(), pair.second->local_path.c_str());
             unique_tasks_vec.push_back(*pair.second);
         }
         common_download_run_tasks(unique_tasks_vec);

ggml/CMakeLists.txt

@@ -342,9 +342,6 @@ set(GGML_PUBLIC_HEADERS
     include/gguf.h)
 
 set_target_properties(ggml PROPERTIES PUBLIC_HEADER "${GGML_PUBLIC_HEADERS}")
-#if (GGML_METAL)
-#    set_target_properties(ggml PROPERTIES RESOURCE "${CMAKE_CURRENT_SOURCE_DIR}/src/ggml-metal.metal")
-#endif()
 install(TARGETS ggml LIBRARY PUBLIC_HEADER)
 install(TARGETS ggml-base LIBRARY)
 

ggml/include/ggml-rpc.h

@@ -8,10 +8,10 @@ extern "C" {
 
 #define RPC_PROTO_MAJOR_VERSION    4
 #define RPC_PROTO_MINOR_VERSION    0
-#define RPC_PROTO_PATCH_VERSION    2
+#define RPC_PROTO_PATCH_VERSION    1
 
 #ifdef  __cplusplus
-static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
 #endif
 
 #define GGML_RPC_MAX_SERVERS       16

ggml/include/ggml.h

@@ -570,7 +570,6 @@ extern "C" {
         GGML_OP_RWKV_WKV7,
         GGML_OP_SOLVE_TRI,
         GGML_OP_GATED_DELTA_NET,
-        GGML_OP_LIGHTNING_INDEXER,
 
         GGML_OP_UNARY,
 
@@ -2576,24 +2575,6 @@ extern "C" {
             struct ggml_tensor  * state,
             int64_t               K);
 
-    // DSA lightning indexer
-    //
-    // q:       [n_embd_idx, n_head_idx, n_batch, ne3 ]
-    // k:       [n_embd_idx, 1,          n_kv,    ne3 ]
-    // weights: [n_head_idx, n_batch,    1,       ne3 ] !! prescaled !!
-    // mask:    [n_kv,       n_batch,    1,       ne33] !! f16 !!
-    // res:     [n_kv,       n_batch,    1,       ne3 ]
-    //
-    // broadcast:
-    //   ne3 % ne33 == 0
-    //
-    GGML_API struct ggml_tensor * ggml_lightning_indexer(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * q,
-        struct ggml_tensor  * k,
-        struct ggml_tensor  * weights,
-        struct ggml_tensor  * mask);
-
     // custom operators
 
     typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata);

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

@@ -2060,10 +2060,6 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_gated_delta_net(params, tensor);
             } break;
-        case GGML_OP_LIGHTNING_INDEXER:
-            {
-                ggml_compute_forward_lightning_indexer(params, tensor);
-            } break;
         case GGML_OP_MAP_CUSTOM1:
             {
                 ggml_compute_forward_map_custom1(params, tensor);
@@ -2384,7 +2380,6 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_FLASH_ATTN_BACK:
         case GGML_OP_SSM_CONV:
         case GGML_OP_SSM_SCAN:
-        case GGML_OP_LIGHTNING_INDEXER:
             {
                 n_tasks = n_threads;
             } break;
@@ -2970,12 +2965,6 @@ struct ggml_cplan ggml_graph_plan(
                     {
                         GGML_ABORT("fatal error");
                     }
-                case GGML_OP_LIGHTNING_INDEXER:
-                    {
-                        // temp buffer for dequantizing lightning indexer keys
-                        const int64_t ne10 = node->src[1]->ne[0];
-                        cur += sizeof(float)*ne10*n_tasks;
-                    } break;
                 default:
                     break;
             }

ggml/src/ggml-cpu/ops.cpp

@@ -11568,87 +11568,3 @@ void ggml_compute_forward_fwht(const ggml_compute_params * params, ggml_tensor *
             }
     }
 }
-
-// ggml_compute_forward_lightning_indexer
-
-void ggml_compute_forward_lightning_indexer(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-
-    const ggml_tensor * q = dst->src[0];
-    const ggml_tensor * k = dst->src[1];
-    const ggml_tensor * w = dst->src[2]; // weights
-    const ggml_tensor * m = dst->src[3]; // mask
-
-    GGML_ASSERT(dst->type  == GGML_TYPE_F32);
-    GGML_ASSERT(   q->type == GGML_TYPE_F32);
-    GGML_ASSERT(   w->type == GGML_TYPE_F32);
-    GGML_ASSERT(   m->type == GGML_TYPE_F16);
-
-    GGML_TENSOR_LOCALS(int64_t, neq,  q, ne)
-    GGML_TENSOR_LOCALS(size_t,  nbq,  q, nb)
-    GGML_TENSOR_LOCALS(int64_t, nek,  k, ne)
-    GGML_TENSOR_LOCALS(size_t,  nbk,  k, nb)
-    GGML_TENSOR_LOCALS(int64_t, new,  w, ne)
-    GGML_TENSOR_LOCALS(size_t,  nbw,  w, nb)
-    GGML_TENSOR_LOCALS(int64_t, nem,  m, ne)
-    GGML_TENSOR_LOCALS(size_t,  nbm,  m, nb)
-    GGML_TENSOR_LOCALS(int64_t, ne, dst, ne)
-    GGML_TENSOR_LOCALS(size_t,  nb, dst, nb)
-
-    GGML_ASSERT( nb0 == ggml_type_size(dst->type));
-    GGML_ASSERT(nbq0 == ggml_type_size(  q->type));
-    GGML_ASSERT(nbk0 == ggml_type_size(  k->type));
-    GGML_ASSERT(nbw0 == ggml_type_size(  w->type));
-    GGML_ASSERT(nbm0 == ggml_type_size(  m->type));
-
-    const int n_embd    = q->ne[0];
-    const int n_head    = q->ne[1];
-    const int n_tokens  = q->ne[2];
-    const int n_stream  = q->ne[3];
-    const int n_kv      = k->ne[2];
-
-    ggml_to_float_t const k_to_float = ggml_get_type_traits(k->type)->to_float;
-    GGML_ASSERT((k->type == GGML_TYPE_F32 || k_to_float) && "lightning indexer: unsupported K-type");
-
-    const int nr  = n_kv;
-    const int ith = params->ith;
-    const int nth = params->nth;
-
-    // (temporary) buffer for K converted to float
-    float * k_row_f32 = (float *) params->wdata + ith*(1*n_embd + CACHE_LINE_SIZE_F32);
-
-    // rows per thread
-    const int dr = (nr + nth - 1)/nth;
-
-    // row range for this thread
-    const int ir0 = dr*ith;
-    const int ir1 = MIN(ir0 + dr, nr);
-
-    for (int s = 0; s < n_stream; ++s) {
-        for (int t = 0; t < n_tokens; ++t) {
-            const float       *   w_row =       (float *) ((char *)   w->data + t*nbw1 +        s*nbw3);
-            const ggml_fp16_t *   m_row = (ggml_fp16_t *) ((char *)   m->data + t*nbm1 + (s%nem3)*nbm3);
-            float             * dst_row =       (float *) ((char *) dst->data + t*nb1  +        s*nb3 );
-            for (int ik = ir0; ik < ir1; ++ik) {
-                char * k_row = (char *) k->data + ik*nbk2 + s*nbk3;
-                if (k_to_float) {
-                    k_to_float(k_row, k_row_f32, n_embd);
-                } else {
-                    k_row_f32 = (float *) k_row;
-                }
-                float score = 0.0f;
-                for (int h = 0; h < n_head; ++h) {
-                    // dot product of q and k for head h
-                    float qk = 0.0f;
-                    const float * q_row = (float *) ((char *) q->data + h*nbq1 + t*nbq2 + s*nbq3);
-                    ggml_vec_dot_f32(n_embd, &qk, 0, q_row, 0, k_row_f32, 0, 1);
-                    // ReLU and weights (prescaled)
-                    score += MAX(qk, 0.0f) * w_row[h];
-                }
-                // apply mask
-                dst_row[ik] = score + GGML_CPU_FP16_TO_FP32(m_row[ik]);
-            }
-        }
-    }
-}

ggml/src/ggml-cpu/ops.h

@@ -105,7 +105,6 @@ void ggml_compute_forward_rwkv_wkv7(const struct ggml_compute_params * params, s
 void ggml_compute_forward_solve_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_gla(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_gated_delta_net(const struct ggml_compute_params * params, struct ggml_tensor * dst);
-void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom1(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom2(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom3(const struct ggml_compute_params * params, struct ggml_tensor * dst);

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

@@ -4493,14 +4493,7 @@ static bool ggml_backend_cuda_get_available_uma_memory(long * available_memory_k
 static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) {
     ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context;
     ggml_cuda_set_device(ctx->device);
-    cudaError_t err = cudaMemGetInfo(free, total);
-    if (err != cudaSuccess) {
-        (void)cudaGetLastError();
-        GGML_LOG_WARN("%s: cudaMemGetInfo failed (%s), returning 0/0\n", __func__, cudaGetErrorString(err));
-        *free = 0;
-        *total = 0;
-        return;
-    }
+    CUDA_CHECK(cudaMemGetInfo(free, total));
 
 // ref: https://github.com/ggml-org/llama.cpp/pull/17368
 #if defined(__linux__)

ggml/src/ggml-hexagon/htp/CMakeLists.txt

@@ -31,6 +31,7 @@ add_library(${HTP_LIB} SHARED
     get-rows-ops.c
     cpy-ops.c
     repeat-ops.c
+    argsort-ops.c
     ssm-conv.c
     cumsum-ops.c
     fill-ops.c
@@ -38,9 +39,8 @@ add_library(${HTP_LIB} SHARED
     diag-ops.c
     solve-tri-ops.c
     pad-ops.c
-    flash-attn-ops.c
     matmul-ops.c
-    argsort-ops.c
+    flash-attn-ops.c
 )
 
 target_compile_definitions(${HTP_LIB} PRIVATE

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

@@ -22,8 +22,6 @@
 struct htp_argsort_context {
     struct htp_ops_context * octx;
     uint32_t                 nrows_per_thread;
-    uint8_t *                vtcm_base;
-    size_t                   vtcm_per_thread;
 };
 
 static inline bool all_greater_f32(HVX_Vector x, HVX_Vector y)
@@ -172,208 +170,7 @@ int32_t argosrt_ramp_lut[32] __attribute__((aligned(VLEN))) = {
     16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31
 };
 
-__attribute__((always_inline))
-static inline void vec_cas(HVX_Vector * X_val, HVX_Vector * X_idx, HVX_Vector * Y_val, HVX_Vector * Y_idx, bool asc) {
-    HVX_VectorPred pred = asc ? Q6_Q_vcmp_gt_VsfVsf(*X_val, *Y_val)
-                              : Q6_Q_vcmp_gt_VsfVsf(*Y_val, *X_val);
-    HVX_Vector next_X_val = Q6_V_vmux_QVV(pred, *Y_val, *X_val);
-    HVX_Vector next_Y_val = Q6_V_vmux_QVV(pred, *X_val, *Y_val);
-    HVX_Vector next_X_idx = Q6_V_vmux_QVV(pred, *Y_idx, *X_idx);
-    HVX_Vector Y_tmp_idx  = Q6_V_vmux_QVV(pred, *X_idx, *Y_idx);
-    *X_val = next_X_val;
-    *Y_val = next_Y_val;
-    *X_idx = next_X_idx;
-    *Y_idx = Y_tmp_idx;
-}
-
-__attribute__((always_inline))
-static inline void bitonic_cas_32(HVX_Vector * V, HVX_Vector * I, int d, HVX_VectorPred dir_mask, HVX_Vector idx_vec, HVX_Vector zero_vec) {
-    HVX_VectorPred mask_left;
-    HVX_Vector V_rot_left, V_rot_right;
-    HVX_Vector I_rot_left, I_rot_right;
-
-    if (d == 1) {
-        mask_left = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(1)), zero_vec);
-        V_rot_left = Q6_V_vror_VR(*V, 4);
-        V_rot_right = Q6_V_vror_VR(*V, 124);
-        I_rot_left = Q6_V_vror_VR(*I, 4);
-        I_rot_right = Q6_V_vror_VR(*I, 124);
-    } else if (d == 2) {
-        mask_left = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(2)), zero_vec);
-        V_rot_left = Q6_V_vror_VR(*V, 8);
-        V_rot_right = Q6_V_vror_VR(*V, 120);
-        I_rot_left = Q6_V_vror_VR(*I, 8);
-        I_rot_right = Q6_V_vror_VR(*I, 120);
-    } else if (d == 4) {
-        mask_left = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(4)), zero_vec);
-        V_rot_left = Q6_V_vror_VR(*V, 16);
-        V_rot_right = Q6_V_vror_VR(*V, 112);
-        I_rot_left = Q6_V_vror_VR(*I, 16);
-        I_rot_right = Q6_V_vror_VR(*I, 112);
-    } else if (d == 8) {
-        mask_left = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(8)), zero_vec);
-        V_rot_left = Q6_V_vror_VR(*V, 32);
-        V_rot_right = Q6_V_vror_VR(*V, 96);
-        I_rot_left = Q6_V_vror_VR(*I, 32);
-        I_rot_right = Q6_V_vror_VR(*I, 96);
-    } else { // d == 16
-        mask_left = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(16)), zero_vec);
-        V_rot_left = Q6_V_vror_VR(*V, 64);
-        V_rot_right = Q6_V_vror_VR(*V, 64);
-        I_rot_left = Q6_V_vror_VR(*I, 64);
-        I_rot_right = Q6_V_vror_VR(*I, 64);
-    }
-
-    HVX_Vector V_paired = Q6_V_vmux_QVV(mask_left, V_rot_left, V_rot_right);
-    HVX_Vector I_paired = Q6_V_vmux_QVV(mask_left, I_rot_left, I_rot_right);
-
-    HVX_VectorPred V_gt_Vpaired = Q6_Q_vcmp_gt_VsfVsf(*V, V_paired);
-    HVX_VectorPred Vpaired_gt_V = Q6_Q_vcmp_gt_VsfVsf(V_paired, *V);
-    HVX_VectorPred mask_right = Q6_Q_not_Q(mask_left);
-    HVX_VectorPred Q_asc = Q6_Q_or_QQ(
-        Q6_Q_and_QQ(mask_left, V_gt_Vpaired),
-        Q6_Q_and_QQ(Vpaired_gt_V, mask_right)
-    );
-    HVX_VectorPred Q_swap = Q6_Q_or_QQ(
-        Q6_Q_and_QQ(dir_mask, Q_asc),
-        Q6_Q_and_QQ(Q6_Q_not_Q(dir_mask), Q6_Q_not_Q(Q_asc))
-    );
-
-    *V = Q6_V_vmux_QVV(Q_swap, V_paired, *V);
-    *I = Q6_V_vmux_QVV(Q_swap, I_paired, *I);
-}
-
-__attribute__((always_inline))
-static inline void bitonic_sort_generic_hvx(uint8_t * values, uint8_t * indices, int K, bool asc_order) {
-    HVX_Vector V[32];
-    HVX_Vector I[32];
-
-    HVX_Vector zero_vec = Q6_V_vzero();
-    HVX_Vector idx_vec = *(HVX_Vector *)argosrt_ramp_lut;
-
-    // Load values and initialize indices
-    for (int v = 0; v < K; v++) {
-        V[v] = *(HVX_Vector *)(values + v * 128);
-        I[v] = Q6_Vw_vadd_VwVw(idx_vec, Q6_V_vsplat_R(v * 32));
-    }
-
-    HVX_VectorPred pred_all_1s = Q6_Q_vcmp_eq_VwVw(zero_vec, zero_vec);
-    HVX_VectorPred pred_all_0s = Q6_Q_not_Q(pred_all_1s);
-
-    int M = 5;
-    while ((1 << (M - 5)) < K) M++;
-
-    for (int s = 1; s <= M; s++) {
-        for (int stage_d = s - 1; stage_d >= 0; stage_d--) {
-            int d = 1 << stage_d;
-            if (d >= 32) {
-                int v_dist = d / 32;
-                for (int v1 = 0; v1 < K; v1++) {
-                    if ((v1 & v_dist) == 0) {
-                        int v2 = v1 + v_dist;
-                        bool asc = (s < M) ? ((((v1 * 32) >> s) % 2) == 0) : asc_order;
-                        vec_cas(&V[v1], &I[v1], &V[v2], &I[v2], asc);
-                    }
-                }
-            } else {
-                if (s < 5) {
-                    HVX_VectorPred dir_mask = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(1 << s)), zero_vec);
-                    for (int v = 0; v < K; v++) {
-                        bitonic_cas_32(&V[v], &I[v], d, dir_mask, idx_vec, zero_vec);
-                    }
-                } else {
-                    for (int v = 0; v < K; v++) {
-                        bool asc = (s < M) ? ((((v * 32) >> s) % 2) == 0) : asc_order;
-                        HVX_VectorPred dir_mask = asc ? pred_all_1s : pred_all_0s;
-                        bitonic_cas_32(&V[v], &I[v], d, dir_mask, idx_vec, zero_vec);
-                    }
-                }
-            }
-        }
-    }
-
-    // Write back sorted values and indices
-    for (int v = 0; v < K; v++) {
-        *(HVX_Vector *)(values + v * 128)  = V[v];
-        *(HVX_Vector *)(indices + v * 128) = I[v];
-    }
-}
-
-__attribute__((always_inline))
-static inline void sort32_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
-    bitonic_sort_generic_hvx(values, indices, 1, order == GGML_SORT_ORDER_ASC);
-}
-
-__attribute__((always_inline))
-static inline void sort64_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
-    bitonic_sort_generic_hvx(values, indices, 2, order == GGML_SORT_ORDER_ASC);
-}
-
-__attribute__((always_inline))
-static inline void sort128_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
-    bitonic_sort_generic_hvx(values, indices, 4, order == GGML_SORT_ORDER_ASC);
-}
-
-__attribute__((always_inline))
-static inline void sort256_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
-    bitonic_sort_generic_hvx(values, indices, 8, order == GGML_SORT_ORDER_ASC);
-}
-
-__attribute__((always_inline))
-static inline void sort512_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
-    bitonic_sort_generic_hvx(values, indices, 16, order == GGML_SORT_ORDER_ASC);
-}
-
-__attribute__((always_inline))
-static inline void sort1024_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
-    bitonic_sort_generic_hvx(values, indices, 32, order == GGML_SORT_ORDER_ASC);
-}
-
-#define HTP_ARGSORT_FN(ne00, order_name, order_enum, sort_fn)                                                  \
-static void htp_argsort_f32_##ne00##_##order_name(unsigned int n, unsigned int i, void * data) {               \
-    struct htp_argsort_context * actx = (struct htp_argsort_context *)data;                                    \
-    struct htp_ops_context * octx = actx->octx;                                                                \
-    const struct htp_tensor * src0 = octx->src[0];                                                             \
-    const struct htp_tensor * dst = octx->dst;                                                                 \
-    uint8_t * spad = actx->vtcm_base + actx->vtcm_per_thread * i;                                              \
-    uint32_t total_rows = src0->ne[1] * src0->ne[2] * src0->ne[3];                                             \
-    uint32_t rows_per_thread = actx->nrows_per_thread;                                                         \
-    uint32_t start_row = rows_per_thread * i;                                                                  \
-    uint32_t end_row = MIN(start_row + rows_per_thread, total_rows);                                           \
-    size_t values_size = hex_round_up(ne00 * sizeof(float), 128);                                              \
-    float * values_buf = (float *) spad;                                                                       \
-    int32_t * indices_buf = (int32_t *) (spad + values_size);                                                  \
-    uint32_t nb01 = src0->nb[1];                                                                               \
-    uint32_t nb1 = dst->nb[1];                                                                                 \
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[i] : NULL;                                    \
-    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, start_row);                                              \
-    for (uint32_t r = start_row; r < end_row; r++) {                                                           \
-        uint32_t src_offset = r * nb01;                                                                        \
-        uint32_t dst_offset = r * nb1;                                                                         \
-        uint8_t * src_ptr = (uint8_t *) src0->data + src_offset;                                               \
-        uint8_t * dst_ptr = (uint8_t *) dst->data  + dst_offset;                                               \
-        hex_l2fetch(src_ptr, ne00 * sizeof(float), ne00 * sizeof(float), 1);                                   \
-        hvx_copy_f32_au((uint8_t*)values_buf, src_ptr, ne00);                                                  \
-        sort_fn((uint8_t*)values_buf, (uint8_t*)indices_buf, order_enum);                                      \
-        hvx_copy_f32_ua(dst_ptr, (const uint8_t *) indices_buf, ne00);                                         \
-    }                                                                                                          \
-    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, start_row);                                               \
-}
-
-HTP_ARGSORT_FN(32,   asc, GGML_SORT_ORDER_ASC,  sort32_f32_hvx)
-HTP_ARGSORT_FN(32,   dsc, GGML_SORT_ORDER_DESC, sort32_f32_hvx)
-HTP_ARGSORT_FN(64,   asc, GGML_SORT_ORDER_ASC,  sort64_f32_hvx)
-HTP_ARGSORT_FN(64,   dsc, GGML_SORT_ORDER_DESC, sort64_f32_hvx)
-HTP_ARGSORT_FN(128,  asc, GGML_SORT_ORDER_ASC,  sort128_f32_hvx)
-HTP_ARGSORT_FN(128,  dsc, GGML_SORT_ORDER_DESC, sort128_f32_hvx)
-HTP_ARGSORT_FN(256,  asc, GGML_SORT_ORDER_ASC,  sort256_f32_hvx)
-HTP_ARGSORT_FN(256,  dsc, GGML_SORT_ORDER_DESC, sort256_f32_hvx)
-HTP_ARGSORT_FN(512,  asc, GGML_SORT_ORDER_ASC,  sort512_f32_hvx)
-HTP_ARGSORT_FN(512,  dsc, GGML_SORT_ORDER_DESC, sort512_f32_hvx)
-HTP_ARGSORT_FN(1024, asc, GGML_SORT_ORDER_ASC,  sort1024_f32_hvx)
-HTP_ARGSORT_FN(1024, dsc, GGML_SORT_ORDER_DESC, sort1024_f32_hvx)
-
-static void htp_argsort_f32_fallback(unsigned int n, unsigned int i, void * data) {
+static void htp_argsort_f32(unsigned int n, unsigned int i, void * data) {
     struct htp_argsort_context * actx = (struct htp_argsort_context *)data;
     struct htp_ops_context * octx = actx->octx;
 
@@ -382,7 +179,7 @@ static void htp_argsort_f32_fallback(unsigned int n, unsigned int i, void * data
     const struct htp_tensor * dst = octx->dst;
 
     // Scratchpad memory
-    uint8_t * spad = actx->vtcm_base + actx->vtcm_per_thread * i;
+    uint8_t * spad = octx->src0_spad.data + octx->src0_spad.size_per_thread * i;
 
     // Dimensions
     uint32_t ne00 = src0->ne[0];
@@ -391,8 +188,12 @@ static void htp_argsort_f32_fallback(unsigned int n, unsigned int i, void * data
     uint32_t ne03 = src0->ne[3];
 
     uint32_t nb01 = src0->nb[1];
+    //uint32_t nb02 = src0->nb[2];
+    //uint32_t nb03 = src0->nb[3];
 
     uint32_t nb1 = dst->nb[1];
+    //uint32_t nb2 = dst->nb[2];
+    //uint32_t nb3 = dst->nb[3];
 
     // Sort order
     enum ggml_sort_order order = (enum ggml_sort_order) octx->op_params[0];
@@ -403,17 +204,20 @@ static void htp_argsort_f32_fallback(unsigned int n, unsigned int i, void * data
     uint32_t start_row = rows_per_thread * i;
     uint32_t end_row = MIN(start_row + rows_per_thread, total_rows);
 
+    // Scratchpad layout:
+    // We need space for one row of float data (values) and one row of int32 indices.
+    // values: ne00 * sizeof(float)
+    // indices: ne00 * sizeof(int32_t)
+    // Padded to 128 bytes.
+
     size_t values_size = hex_round_up(ne00 * sizeof(float), 128);
-    uint32_t num_vec_ind_values = hmx_ceil_div(ne00, VLEN/(sizeof(int32_t)));
+    size_t num_vec_ind_values = hmx_ceil_div(ne00, VLEN/(sizeof(int32_t)));
     float * values_buf = (float *) spad;
     int32_t * indices_buf = (int32_t *) (spad + values_size);
     HVX_Vector * indices_buf_vec = (HVX_Vector *) (spad + values_size);
     const HVX_Vector ind_init_vec = *(HVX_Vector *)argosrt_ramp_lut;
     const HVX_Vector ind_diff_vec = Q6_V_vsplat_R(32);
 
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[i] : NULL;
-    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, start_row);
-
     for (uint32_t r = start_row; r < end_row; r++) {
         uint32_t src_offset = r * nb01;
         uint32_t dst_offset = r * nb1;
@@ -441,8 +245,6 @@ static void htp_argsort_f32_fallback(unsigned int n, unsigned int i, void * data
         // Copy indices back to DDR
         hvx_copy_f32_ua(dst_ptr, (const uint8_t *) indices_buf, ne00);
     }
-
-    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, start_row);
 }
 
 int op_argsort(struct htp_ops_context * octx) {
@@ -471,6 +273,11 @@ int op_argsort(struct htp_ops_context * octx) {
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
+    octx->src0_spad.data = octx->ctx->vtcm_base;
+    octx->src0_spad.size = total_spad_size;
+    octx->src0_spad.size_per_thread = spad_per_thread;
+    octx->src0_spad.src  = NULL;
+
     FARF(HIGH, "argsort: %ux%ux%ux%u -> %ux%ux%ux%u (0x%x, 0x%x)",
          octx->src[0]->ne[0], octx->src[0]->ne[1], octx->src[0]->ne[2], octx->src[0]->ne[3],
          octx->dst->ne[0], octx->dst->ne[1], octx->dst->ne[2], octx->dst->ne[3],
@@ -479,36 +286,9 @@ int op_argsort(struct htp_ops_context * octx) {
     struct htp_argsort_context actx;
     actx.octx = octx;
     actx.nrows_per_thread = (total_rows + n_threads - 1) / n_threads;
-    actx.vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
-    actx.vtcm_per_thread = spad_per_thread;
-
-    enum ggml_sort_order order = (enum ggml_sort_order) octx->op_params[0];
-    worker_callback_t job_func = htp_argsort_f32_fallback;
-
-    if (order == GGML_SORT_ORDER_ASC) {
-        switch (ne00) {
-            case 1024: job_func = htp_argsort_f32_1024_asc; break;
-            case 512:  job_func = htp_argsort_f32_512_asc;  break;
-            case 256:  job_func = htp_argsort_f32_256_asc;  break;
-            case 128:  job_func = htp_argsort_f32_128_asc;  break;
-            case 64:   job_func = htp_argsort_f32_64_asc;   break;
-            case 32:   job_func = htp_argsort_f32_32_asc;   break;
-            default:   job_func = htp_argsort_f32_fallback; break;
-        }
-    } else {
-        switch (ne00) {
-            case 1024: job_func = htp_argsort_f32_1024_dsc; break;
-            case 512:  job_func = htp_argsort_f32_512_dsc;  break;
-            case 256:  job_func = htp_argsort_f32_256_dsc;  break;
-            case 128:  job_func = htp_argsort_f32_128_dsc;  break;
-            case 64:   job_func = htp_argsort_f32_64_dsc;   break;
-            case 32:   job_func = htp_argsort_f32_32_dsc;   break;
-            default:   job_func = htp_argsort_f32_fallback; break;
-        }
-    }
 
     // Run jobs
-    worker_pool_run_func(octx->ctx->worker_pool, job_func, &actx, n_threads);
+    worker_pool_run_func(octx->ctx->worker_pool, htp_argsort_f32, &actx, n_threads);
 
     return HTP_STATUS_OK;
 }

ggml/src/ggml-metal/CMakeLists.txt

@@ -24,62 +24,119 @@ if (GGML_METAL_NDEBUG)
 endif()
 
 set(METALLIB_COMMON "${CMAKE_CURRENT_SOURCE_DIR}/../ggml-common.h")
+set(METALLIB_KERNELS_COMMON     "${CMAKE_CURRENT_SOURCE_DIR}/kernels/common.h")
+set(METALLIB_KERNELS_DEQUANTIZE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/dequantize.h")
+set(METALLIB_KERNELS_QUANTIZE   "${CMAKE_CURRENT_SOURCE_DIR}/kernels/quantize.h")
+
+set(METALLIB_KERNEL_SOURCES
+    kernels/fa.metal
+    kernels/mul_mv.metal
+    kernels/mul_mm.metal
+    kernels/quantize.metal
+    kernels/softmax.metal
+    kernels/norm.metal
+    kernels/unary.metal
+    kernels/binbcast.metal
+    kernels/reduce.metal
+    kernels/tri.metal
+    kernels/ssm.metal
+    kernels/wkv.metal
+    kernels/gated_delta_net.metal
+    kernels/solve_tri.metal
+    kernels/rope.metal
+    kernels/conv.metal
+    kernels/upscale.metal
+    kernels/argsort.metal
+    kernels/pool.metal
+    kernels/misc.metal
+)
+
 if (GGML_METAL_EMBED_LIBRARY)
     enable_language(ASM)
 
     add_compile_definitions(GGML_METAL_EMBED_LIBRARY)
 
-    set(METALLIB_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/ggml-metal.metal")
-    set(METALLIB_IMPL   "${CMAKE_CURRENT_SOURCE_DIR}/ggml-metal-impl.h")
+    set(METALLIB_IMPL "${CMAKE_CURRENT_SOURCE_DIR}/ggml-metal-impl.h")
 
     file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/autogenerated")
 
-    # merge ggml-common.h and ggml-metal.metal into a single file
-    set(METALLIB_EMBED_ASM        "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed.s")
-    set(METALLIB_SOURCE_EMBED     "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed.metal")
-    set(METALLIB_SOURCE_EMBED_TMP "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed.metal.tmp")
+    set(METALLIB_EMBED_ASM_FILES "")
+    foreach(src ${METALLIB_KERNEL_SOURCES})
+        get_filename_component(kind ${src} NAME_WE)
+        # symbol names must be valid C identifiers ('-' is not allowed)
+        string(REPLACE "-" "_" kind_sym ${kind})
 
-    add_custom_command(
-        OUTPUT "${METALLIB_EMBED_ASM}"
-        COMMAND echo "Embedding Metal library"
-        COMMAND sed -e "/__embed_ggml-common.h__/r ${METALLIB_COMMON}"       -e "/__embed_ggml-common.h__/d"         < "${METALLIB_SOURCE}"           > "${METALLIB_SOURCE_EMBED_TMP}"
-        COMMAND sed -e "/\#include \"ggml-metal-impl.h\"/r ${METALLIB_IMPL}" -e "/\#include \"ggml-metal-impl.h\"/d" < "${METALLIB_SOURCE_EMBED_TMP}" > "${METALLIB_SOURCE_EMBED}"
-        COMMAND echo ".section __DATA,__ggml_metallib"          >  "${METALLIB_EMBED_ASM}"
-        COMMAND echo ".globl _ggml_metallib_start"              >> "${METALLIB_EMBED_ASM}"
-        COMMAND echo "_ggml_metallib_start:"                    >> "${METALLIB_EMBED_ASM}"
-        COMMAND echo .incbin "\"${METALLIB_SOURCE_EMBED}\""     >> "${METALLIB_EMBED_ASM}"
-        COMMAND echo ".globl _ggml_metallib_end"                >> "${METALLIB_EMBED_ASM}"
-        COMMAND echo "_ggml_metallib_end:"                      >> "${METALLIB_EMBED_ASM}"
-        DEPENDS ../ggml-common.h ggml-metal.metal ggml-metal-impl.h
-        COMMENT "Generate assembly for embedded Metal library"
-        VERBATIM
-    )
+        set(SRC   "${CMAKE_CURRENT_SOURCE_DIR}/kernels/${kind}.metal")
+        set(EMBED "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed-${kind}.metal")
+        set(ASM   "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed-${kind}.s")
 
-    target_sources(ggml-metal PRIVATE "${METALLIB_EMBED_ASM}")
+        # only prepend headers that this source actually includes
+        set(HEADERS_FOR_SRC ${METALLIB_KERNELS_COMMON})
+        file(STRINGS ${SRC} _has_dequantize REGEX "#include \"dequantize\\.h\"")
+        file(STRINGS ${SRC} _has_quantize   REGEX "#include \"quantize\\.h\"")
+        if(_has_dequantize)
+            list(APPEND HEADERS_FOR_SRC ${METALLIB_KERNELS_DEQUANTIZE})
+        endif()
+        if(_has_quantize)
+            list(APPEND HEADERS_FOR_SRC ${METALLIB_KERNELS_QUANTIZE})
+        endif()
+
+        add_custom_command(
+            OUTPUT "${ASM}"
+            # Step 1: concatenate shared headers + this kernel source
+            COMMAND cat ${HEADERS_FOR_SRC} ${SRC} > "${EMBED}.tmp1"
+            # Step 2: remove internal #include and #pragma once
+            COMMAND sed -e "/\#include \"common.h\"/d" -e "/\#include \"dequantize.h\"/d" -e "/\#include \"quantize.h\"/d" -e "/\#pragma once/d" < "${EMBED}.tmp1" > "${EMBED}.tmp2"
+            # Step 3: inline ggml-common.h (replacing __embed_ggml-common.h__ sentinel)
+            COMMAND sed -e "/__embed_ggml-common.h__/r ${METALLIB_COMMON}" -e "/__embed_ggml-common.h__/d" < "${EMBED}.tmp2" > "${EMBED}.tmp3"
+            # Step 4: inline ggml-metal-impl.h
+            COMMAND sed -e "/\#include \"ggml-metal-impl.h\"/r ${METALLIB_IMPL}" -e "/\#include \"ggml-metal-impl.h\"/d" < "${EMBED}.tmp3" > "${EMBED}"
+            # Step 5: emit an asm chunk with kind-specific start/end symbols
+            #   note: '-' is illegal in C symbols, so we use kind_sym; the macOS
+            #   section name is limited to 16 chars so we keep it shared
+            #   across kinds (__ggml_metallib) and only vary the global symbols.
+            COMMAND echo ".section __DATA,__ggml_metallib"                       >  "${ASM}"
+            COMMAND echo ".globl _ggml_metallib_${kind_sym}_start"               >> "${ASM}"
+            COMMAND echo "_ggml_metallib_${kind_sym}_start:"                     >> "${ASM}"
+            COMMAND echo .incbin "\"${EMBED}\""                                  >> "${ASM}"
+            COMMAND echo ".globl _ggml_metallib_${kind_sym}_end"                 >> "${ASM}"
+            COMMAND echo "_ggml_metallib_${kind_sym}_end:"                       >> "${ASM}"
+            DEPENDS ../ggml-common.h ggml-metal-impl.h
+                    kernels/common.h kernels/dequantize.h kernels/quantize.h
+                    kernels/${kind}.metal
+            COMMENT "Generate embedded Metal library for ${kind}"
+            VERBATIM
+        )
+
+        list(APPEND METALLIB_EMBED_ASM_FILES "${ASM}")
+    endforeach()
+
+    target_sources(ggml-metal PRIVATE ${METALLIB_EMBED_ASM_FILES})
 else()
-    # copy metal files to bin directory
+    # copy header files to bin directory
     configure_file(../ggml-common.h  ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-common.h     COPYONLY)
-    configure_file(ggml-metal.metal  ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal  COPYONLY)
     configure_file(ggml-metal-impl.h ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal-impl.h COPYONLY)
 
+    file(MAKE_DIRECTORY "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels")
+    configure_file(kernels/common.h     ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/common.h     COPYONLY)
+    configure_file(kernels/dequantize.h ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/dequantize.h COPYONLY)
+    configure_file(kernels/quantize.h   ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/quantize.h   COPYONLY)
+
+    foreach(src ${METALLIB_KERNEL_SOURCES})
+        configure_file(${src} ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${src} COPYONLY)
+    endforeach()
+
     if (GGML_METAL_SHADER_DEBUG)
-        # custom command to do the following:
-        #   xcrun -sdk macosx metal    -fno-fast-math -c ggml-metal.metal -o ggml-metal.air
-        #   xcrun -sdk macosx metallib                   ggml-metal.air   -o default.metallib
-        #
-        # note: this is the only way I found to disable fast-math in Metal. it's ugly, but at least it works
-        #       disabling fast math is needed in order to pass tests/test-backend-ops
+        # note: disabling fast math is needed in order to pass tests/test-backend-ops
         # note: adding -fno-inline fixes the tests when using MTL_SHADER_VALIDATION=1
         # note: unfortunately, we have to call it default.metallib instead of ggml.metallib
         #       ref: https://github.com/ggml-org/whisper.cpp/issues/1720
         # note: adding -g causes segmentation fault during compile
-        #set(XC_FLAGS -fno-fast-math -fno-inline -g)
         set(XC_FLAGS -fno-fast-math -fno-inline)
     else()
         set(XC_FLAGS -O3)
     endif()
 
-    # Append macOS metal versioning flags
     if (GGML_METAL_MACOSX_VERSION_MIN)
         message(STATUS "Adding  -mmacosx-version-min=${GGML_METAL_MACOSX_VERSION_MIN} flag to metal compilation")
         list   (APPEND XC_FLAGS -mmacosx-version-min=${GGML_METAL_MACOSX_VERSION_MIN})
@@ -90,35 +147,46 @@ else()
         list   (APPEND XC_FLAGS -std=${GGML_METAL_STD})
     endif()
 
+    # Compile each kernel source to .air, then link into default.metallib
+    set(AIR_FILES "")
+    foreach(src ${METALLIB_KERNEL_SOURCES})
+        get_filename_component(name ${src} NAME_WE)
+        set(AIR "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${name}.air")
+        list(APPEND AIR_FILES ${AIR})
+        add_custom_command(
+            OUTPUT ${AIR}
+            COMMAND xcrun -sdk macosx metal ${XC_FLAGS} -I ${CMAKE_RUNTIME_OUTPUT_DIRECTORY} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${src} -o ${AIR}
+            DEPENDS ${src} kernels/common.h kernels/dequantize.h kernels/quantize.h ${METALLIB_COMMON} ggml-metal-impl.h
+            COMMENT "Compiling ${src}"
+            VERBATIM
+        )
+    endforeach()
+
     add_custom_command(
         OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
-        COMMAND xcrun -sdk macosx metal ${XC_FLAGS} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal -o - |
-                xcrun -sdk macosx metallib        - -o ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
+        COMMAND xcrun -sdk macosx metallib ${AIR_FILES} -o ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
         COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-common.h
-        COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal
-        DEPENDS ggml-metal.metal ${METALLIB_COMMON}
-        COMMENT "Compiling Metal kernels"
-        )
+        COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal-impl.h
+        COMMAND rm -rf ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels
+        DEPENDS ${AIR_FILES}
+        COMMENT "Linking Metal kernels into default.metallib"
+    )
 
-    # FIXME: only add to the ggml-metal target?
     add_custom_target(
         ggml-metal-lib ALL
         DEPENDS ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
-        )
+    )
 endif() # GGML_METAL_EMBED_LIBRARY
 
 if (NOT GGML_METAL_EMBED_LIBRARY)
     install(
-        FILES src/ggml-metal/ggml-metal.metal
-        PERMISSIONS
-            OWNER_READ
-            OWNER_WRITE
-            GROUP_READ
-            WORLD_READ
-        DESTINATION ${CMAKE_INSTALL_BINDIR})
+        DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/kernels/
+        DESTINATION ${CMAKE_INSTALL_BINDIR}/kernels
+        FILES_MATCHING PATTERN "*.metal" PATTERN "*.h"
+    )
 
-        install(
-            FILES ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
-            DESTINATION ${CMAKE_INSTALL_BINDIR}
-        )
+    install(
+        FILES ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
+        DESTINATION ${CMAKE_INSTALL_BINDIR}
+    )
 endif()

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

@@ -94,8 +94,63 @@ int ggml_metal_pipeline_max_theads_per_threadgroup(struct ggml_metal_pipeline_wi
     return pipeline.pipeline->obj.maxTotalThreadsPerThreadgroup;
 }
 
+//
+// MTLLibrary collection (one library per op-source, compiled separately)
+//
+
+// Single source of truth for the per-kind metal libraries. The order here
+// defines the enum values and every per-kind table below, so adding a library
+// is a one-line change here (plus adding its source to CMakeLists.txt).
+//   X(suffix, name): name is both the kernels/<name>.metal basename and the
+//   ggml_metallib_<name>_{start,end} embed-symbol stem.
+#define GGML_METAL_LIBS \
+    X(FA,              fa)             \
+    X(MUL_MV,          mul_mv)         \
+    X(MUL_MM,          mul_mm)         \
+    X(QUANTIZE,        quantize)       \
+    X(SOFTMAX,         softmax)        \
+    X(NORM,            norm)           \
+    X(UNARY,           unary)          \
+    X(BINBCAST,        binbcast)       \
+    X(REDUCE,          reduce)         \
+    X(TRI,             tri)            \
+    X(SSM,             ssm)            \
+    X(WKV,             wkv)            \
+    X(GATED_DELTA_NET, gated_delta_net)\
+    X(SOLVE_TRI,       solve_tri)      \
+    X(ROPE,            rope)           \
+    X(CONV,            conv)           \
+    X(UPSCALE,         upscale)        \
+    X(ARGSORT,         argsort)        \
+    X(POOL,            pool)           \
+    X(MISC,            misc)
+
+enum ggml_metal_lib_kind {
+#define X(e, s) GGML_METAL_LIB_##e,
+    GGML_METAL_LIBS
+#undef X
+    GGML_METAL_LIB_COUNT,
+};
+
+static const char * const k_lib_names[GGML_METAL_LIB_COUNT] = {
+#define X(e, s) [GGML_METAL_LIB_##e] = #s,
+    GGML_METAL_LIBS
+#undef X
+};
+
 struct ggml_metal_library {
-    id<MTLLibrary> obj;
+    // Per-kind compiled libraries. When single_library is true, the whole library
+    // (e.g. a pre-compiled default.metallib or a from-source build) lives at
+    // objs[0] and the remaining slots are nil.
+    id<MTLLibrary> objs[GGML_METAL_LIB_COUNT];
+    bool single_library; // true: combined library at objs[0]; false: per-kind libs in objs[*]
+
+    // Routing table: kernel function name -> objs[] index, populated from each
+    // compiled library's -[MTLLibrary functionNames]. The actual compiled
+    // libraries are the single source of truth for which library owns a kernel,
+    // so adding kernels later requires no manual routing maintenance.
+    // nil in single_library mode (everything resolves to objs[0]).
+    NSMutableDictionary<NSString *, NSNumber *> * fn_to_lib;
 
     ggml_metal_device_t dev;
     ggml_metal_pipelines_t pipelines; // cache of compiled pipelines
@@ -103,160 +158,376 @@ struct ggml_metal_library {
     NSLock * lock;
 };
 
-ggml_metal_library_t ggml_metal_library_init(ggml_metal_device_t dev) {
-    id<MTLLibrary> library = nil;
-    id<MTLDevice> device = ggml_metal_device_get_obj(dev);
+// Build the fn_to_lib routing table by querying each compiled library's public
+// function names. Call once after all per-kind libraries have been compiled.
+static void ggml_metal_library_build_index(ggml_metal_library_t lib) {
+    @autoreleasepool {
+        NSMutableDictionary<NSString *, NSNumber *> * index = [[NSMutableDictionary alloc] init];
+        for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
+            for (NSString * fname in [lib->objs[kind] functionNames]) {
+                index[fname] = @(kind);
+            }
+        }
+        lib->fn_to_lib = index;
+    }
+}
 
-    // load library
-    //
-    // - first check if the library is embedded
-    // - then check if the library is in the bundle
-    // - if not found, load the source and compile it
-    // - if that fails, return NULL
-    //
-    // TODO: move to a function
-    {
-        const int64_t t_start = ggml_time_us();
+// Parse a `#include "name"` line. Returns the quoted name in *include_name on
+// success. Whitespace-tolerant; ignores `#include <...>` (system headers).
+static bool ggml_metal_library_parse_quoted_include(NSString * line, NSString ** include_name) {
+    NSScanner * scanner = [NSScanner scannerWithString:line];
+    scanner.charactersToBeSkipped = [NSCharacterSet whitespaceCharacterSet];
 
-        NSError * error = nil;
-        NSString * src = nil;
+    if (![scanner scanString:@"#" intoString:NULL] ||
+        ![scanner scanString:@"include" intoString:NULL] ||
+        ![scanner scanString:@"\"" intoString:NULL]) {
+        return false;
+    }
 
-#if GGML_METAL_EMBED_LIBRARY
-        GGML_LOG_INFO("%s: using embedded metal library\n", __func__);
+    NSString * name = nil;
+    if (![scanner scanUpToString:@"\"" intoString:&name]) {
+        return false;
+    }
 
-        extern const char ggml_metallib_start[];
-        extern const char ggml_metallib_end[];
+    if (include_name) {
+        *include_name = name;
+    }
+    return true;
+}
 
-        src = [[NSString alloc] initWithBytes:ggml_metallib_start length:(ggml_metallib_end-ggml_metallib_start) encoding:NSUTF8StringEncoding];
-#else
+// Recursively inline `#include "name"` directives. System includes (<...>),
+// `#if/#else/#endif`, and other preprocessor lines are passed through to the
+// Metal compiler unchanged. `#pragma once` is dropped since `seen` already
+// guards against double-inclusion.
+static bool ggml_metal_library_flatten_file(NSMutableString * dst, NSString * path,
+                                            NSArray<NSString *> * search_paths,
+                                            NSMutableSet<NSString *> * seen, NSError ** error) {
+    NSString * key = [path stringByStandardizingPath];
+    if ([seen containsObject:key]) {
+        return true;
+    }
+    [seen addObject:key];
 
-#ifdef SWIFT_PACKAGE
-        NSBundle * bundle = SWIFTPM_MODULE_BUNDLE;
-#else
-        NSBundle * bundle = [NSBundle bundleForClass:[GGMLMetalClass class]];
-#endif
+    NSString * src = [NSString stringWithContentsOfFile:path encoding:NSUTF8StringEncoding error:error];
+    if (!src) {
+        return false;
+    }
 
-        NSString * path_lib = [bundle pathForResource:@"default" ofType:@"metallib"];
-        if (path_lib == nil) {
-            // Try to find the resource in the directory where the current binary located.
-            NSString * bin_cur = [[NSProcessInfo processInfo] arguments][0];
-            NSString * bin_dir = [bin_cur stringByDeletingLastPathComponent];
+    NSFileManager * fm = [NSFileManager defaultManager];
+    for (NSString * line in [src componentsSeparatedByString:@"\n"]) {
+        NSString * trimmed = [line stringByTrimmingCharactersInSet:[NSCharacterSet whitespaceCharacterSet]];
+        if ([trimmed isEqualToString:@"#pragma once"]) {
+            continue;
+        }
 
-            NSString * path_lib_default = [NSString pathWithComponents:@[bin_dir, @"default.metallib"]];
-            if ([[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
-                GGML_LOG_INFO("%s: found '%s'\n", __func__, [path_lib_default UTF8String]);
-
-                NSDictionary * atts = [[NSFileManager defaultManager] attributesOfItemAtPath:path_lib_default error:&error];
-                if (atts && atts[NSFileType] == NSFileTypeSymbolicLink) {
-                    // Optionally, if this is a symlink, try to resolve it.
-                    path_lib_default = [[NSFileManager defaultManager] destinationOfSymbolicLinkAtPath:path_lib_default error:&error];
-                    if (path_lib_default && [path_lib_default length] > 0 && ![[path_lib_default substringToIndex:1] isEqualToString:@"/"]) {
-                        // It is a relative path, adding the binary directory as directory prefix.
-                        path_lib_default = [NSString pathWithComponents:@[bin_dir, path_lib_default]];
-                    }
-                    if (!path_lib_default || ![[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
-                        // Link to the resource could not be resolved.
-                        path_lib_default = nil;
-                    } else {
-                        GGML_LOG_INFO("%s: symlink resolved '%s'\n", __func__, [path_lib_default UTF8String]);
-                    }
+        NSString * include_name = nil;
+        if (ggml_metal_library_parse_quoted_include(line, &include_name)) {
+            NSString * resolved = nil;
+            for (NSString * dir in search_paths) {
+                NSString * candidate = [dir stringByAppendingPathComponent:include_name];
+                if ([fm isReadableFileAtPath:candidate]) {
+                    resolved = candidate;
+                    break;
                 }
-            } else {
-                // The resource couldn't be found in the binary's directory.
-                path_lib_default = nil;
             }
-
-            path_lib = path_lib_default;
+            if (!resolved) {
+                if (error) {
+                    NSString * msg = [NSString stringWithFormat:@"could not resolve include \"%@\" from '%@'", include_name, path];
+                    *error = [NSError errorWithDomain:@"ggml-metal-source-flatten" code:1
+                                             userInfo:@{NSLocalizedDescriptionKey: msg}];
+                }
+                return false;
+            }
+            if (!ggml_metal_library_flatten_file(dst, resolved, search_paths, seen, error)) {
+                return false;
+            }
+            continue;
         }
 
-        if (path_lib != nil) {
-            // pre-compiled library found
-            NSURL * libURL = [NSURL fileURLWithPath:path_lib];
-            GGML_LOG_INFO("%s: loading '%s'\n", __func__, [path_lib UTF8String]);
+        [dst appendString:line];
+        [dst appendString:@"\n"];
+    }
 
-            library = [device newLibraryWithURL:libURL error:&error];
-            if (error) {
-                GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
-                return nil;
-            }
-        } else {
-            GGML_LOG_INFO("%s: default.metallib not found, loading from source\n", __func__);
+    return true;
+}
 
-            NSString * path_source;
-            NSString * path_resource = [[NSProcessInfo processInfo].environment objectForKey:@"GGML_METAL_PATH_RESOURCES"];
+static NSString * ggml_metal_library_flatten_source(NSString * path_source, NSError ** error) {
+    // Search paths cover both runtime layout (build/bin/kernels + build/bin)
+    // and source-tree layout (ggml/src/ggml-metal/kernels + ggml/src/ggml-metal + ggml/src).
+    NSString * path_kernels = [path_source stringByDeletingLastPathComponent];
+    NSString * path_base    = [path_kernels stringByDeletingLastPathComponent];
+    NSArray<NSString *> * search_paths = @[
+        path_kernels,
+        path_base,
+        [path_base stringByDeletingLastPathComponent],
+    ];
 
-            GGML_LOG_INFO("%s: GGML_METAL_PATH_RESOURCES = %s\n", __func__, path_resource ? [path_resource UTF8String] : "nil");
+    NSMutableString * src = [[NSMutableString alloc] init];
+    NSMutableSet<NSString *> * seen = [NSMutableSet set];
 
-            if (path_resource) {
-                path_source = [path_resource stringByAppendingPathComponent:@"ggml-metal.metal"];
-            } else {
-                path_source = [bundle pathForResource:@"ggml-metal" ofType:@"metal"];
+    if (!ggml_metal_library_flatten_file(src, path_source, search_paths, seen, error)) {
+        [src release];
+        return nil;
+    }
+    return src;
+}
+
+// Compile all per-kind libraries in parallel. `source_for_kind` returns the MSL
+// source for a kind (the helper takes ownership and releases it), or nil with
+// *err set on failure. On success the objs[] slots are populated and the routing
+// index is built; on any failure every error is logged and false is returned
+// (the caller is responsible for freeing `res`).
+static bool ggml_metal_library_compile_all(
+        ggml_metal_library_t res,
+        id<MTLDevice> device,
+        NSDictionary * prep,
+        NSString * (^source_for_kind)(int kind, NSError ** err),
+        const char * origin) {
+    const int64_t t_start = ggml_time_us();
+
+    int64_t  * t_per_lib   = calloc(GGML_METAL_LIB_COUNT, sizeof(int64_t));
+    NSError ** err_per_lib = calloc(GGML_METAL_LIB_COUNT, sizeof(NSError *));
+    __block atomic_bool any_failure = false;
+
+    dispatch_group_t group = dispatch_group_create();
+    dispatch_queue_t queue = dispatch_get_global_queue(QOS_CLASS_USER_INITIATED, 0);
+
+    for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
+        dispatch_group_async(group, queue, ^{
+
+            const int64_t t0 = ggml_time_us();
+
+            NSError * error = nil;
+
+            NSString * src = source_for_kind(kind, &error);
+            if (!src) {
+                err_per_lib[kind] = [error retain];
+                atomic_store(&any_failure, true);
+                return;
             }
 
-            if (path_source == nil) {
-                GGML_LOG_WARN("%s: error: could not use bundle path to find ggml-metal.metal, falling back to trying cwd\n", __func__);
-                path_source = @"ggml-metal.metal";
-            }
+            id<MTLLibrary> lib = nil;
 
-            GGML_LOG_INFO("%s: loading '%s'\n", __func__, [path_source UTF8String]);
-
-            src = [NSString stringWithContentsOfFile:path_source encoding:NSUTF8StringEncoding error:&error];
-            if (error) {
-                GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
-                return nil;
-            }
-        }
-#endif
-
-        if (!library) {
             @autoreleasepool {
-                // dictionary of preprocessor macros
-                NSMutableDictionary * prep = [NSMutableDictionary dictionary];
-
-                if (ggml_metal_device_get_props(dev)->has_bfloat) {
-                    [prep setObject:@"1" forKey:@"GGML_METAL_HAS_BF16"];
-                }
-
-                if (ggml_metal_device_get_props(dev)->has_tensor) {
-                    [prep setObject:@"1" forKey:@"GGML_METAL_HAS_TENSOR"];
-                }
-
-#if GGML_METAL_EMBED_LIBRARY
-                [prep setObject:@"1" forKey:@"GGML_METAL_EMBED_LIBRARY"];
-#endif
-
                 MTLCompileOptions * options = [MTLCompileOptions new];
                 options.preprocessorMacros = prep;
 
-                //[options setFastMathEnabled:false];
+                lib = [device newLibraryWithSource:src options:options error:&error];
 
-                library = [device newLibraryWithSource:src options:options error:&error];
-                if (error) {
-                    GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
-                    return nil;
-                }
-
-#if !__has_feature(objc_arc)
                 [options release];
-#endif
+
+                // retain the error before the autorelease pool drains it
+                if (!lib) {
+                    err_per_lib[kind] = [error retain];
+                }
+            }
+
+            [src release];
+
+            t_per_lib[kind] = ggml_time_us() - t0;
+
+            if (!lib) {
+                atomic_store(&any_failure, true);
+                return;
+            }
+
+            res->objs[kind] = lib;
+        });
+    }
+    dispatch_group_wait(group, DISPATCH_TIME_FOREVER);
+    dispatch_release(group);
+
+    const bool ok = !atomic_load(&any_failure);
+
+    if (ok) {
+        const int64_t t_total = ggml_time_us() - t_start;
+        int64_t t_max = 0;
+        for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
+            GGML_LOG_DEBUG("%s: compiled '%s' library in %.3f sec\n",
+                           __func__, k_lib_names[kind], t_per_lib[kind] / 1e6);
+            if (t_per_lib[kind] > t_max) t_max = t_per_lib[kind];
+        }
+        GGML_LOG_INFO("%s: loaded %d libraries from %s in %.3f sec (max single = %.3f sec)\n",
+                      __func__, GGML_METAL_LIB_COUNT, origin, t_total / 1e6, t_max / 1e6);
+
+        ggml_metal_library_build_index(res);
+    } else {
+        for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
+            if (err_per_lib[kind]) {
+                GGML_LOG_ERROR("%s: failed to build '%s' library: %s\n", __func__,
+                               k_lib_names[kind], [[err_per_lib[kind] description] UTF8String]);
+                [err_per_lib[kind] release];
             }
         }
-
-#if GGML_METAL_EMBED_LIBRARY
-        [src release];
-#endif // GGML_METAL_EMBED_LIBRARY
-
-        GGML_LOG_INFO("%s: loaded in %.3f sec\n", __func__, (ggml_time_us() - t_start) / 1e6);
     }
 
-    ggml_metal_library_t res = calloc(1, sizeof(struct ggml_metal_library));
+    free(err_per_lib);
+    free(t_per_lib);
 
-    res->obj       = library;
+    return ok;
+}
+
+ggml_metal_library_t ggml_metal_library_init(ggml_metal_device_t dev) {
+    id<MTLDevice> device = ggml_metal_device_get_obj(dev);
+
+    ggml_metal_library_t res = calloc(1, sizeof(struct ggml_metal_library));
     res->dev       = dev;
     res->pipelines = ggml_metal_pipelines_init();
     res->lock      = [NSLock new];
 
+    // shared MTLCompileOptions preprocessor macros (matches the build-time defines)
+    NSMutableDictionary * prep = [NSMutableDictionary dictionary];
+    if (ggml_metal_device_get_props(dev)->has_bfloat) {
+        [prep setObject:@"1" forKey:@"GGML_METAL_HAS_BF16"];
+    }
+    if (ggml_metal_device_get_props(dev)->has_tensor) {
+        [prep setObject:@"1" forKey:@"GGML_METAL_HAS_TENSOR"];
+    }
+#if GGML_METAL_EMBED_LIBRARY
+    [prep setObject:@"1" forKey:@"GGML_METAL_EMBED_LIBRARY"];
+#endif
+
+#if GGML_METAL_EMBED_LIBRARY
+    GGML_LOG_INFO("%s: using embedded metal library\n", __func__);
+
+    // start/end symbols emitted by CMake (see CMakeLists.txt), one pair per kind
+#define X(e, s) extern const char ggml_metallib_##s##_start[]; extern const char ggml_metallib_##s##_end[];
+    GGML_METAL_LIBS
+#undef X
+
+    static const char * const lib_start[GGML_METAL_LIB_COUNT] = {
+#define X(e, s) [GGML_METAL_LIB_##e] = ggml_metallib_##s##_start,
+    GGML_METAL_LIBS
+#undef X
+    };
+    static const char * const lib_end[GGML_METAL_LIB_COUNT] = {
+#define X(e, s) [GGML_METAL_LIB_##e] = ggml_metallib_##s##_end,
+    GGML_METAL_LIBS
+#undef X
+    };
+
+    const bool ok = ggml_metal_library_compile_all(res, device, prep,
+        ^NSString * (int kind, NSError ** err) {
+            (void) err;
+            return [[NSString alloc] initWithBytes:lib_start[kind]
+                                            length:(lib_end[kind] - lib_start[kind])
+                                          encoding:NSUTF8StringEncoding];
+        }, "embedded data");
+
+    if (!ok) {
+        ggml_metal_library_free(res);
+        return NULL;
+    }
+
     return res;
+#else
+#ifdef SWIFT_PACKAGE
+    NSBundle * bundle = SWIFTPM_MODULE_BUNDLE;
+#else
+    NSBundle * bundle = [NSBundle bundleForClass:[GGMLMetalClass class]];
+#endif
+
+    const int64_t t_start = ggml_time_us();
+
+    NSError * error = nil;
+    NSString * path_lib = [bundle pathForResource:@"default" ofType:@"metallib"];
+    if (path_lib == nil) {
+        // Try to find the resource in the directory where the current binary located.
+        NSString * bin_cur = [[NSProcessInfo processInfo] arguments][0];
+        NSString * bin_dir = [bin_cur stringByDeletingLastPathComponent];
+
+        NSString * path_lib_default = [NSString pathWithComponents:@[bin_dir, @"default.metallib"]];
+        if ([[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
+            GGML_LOG_INFO("%s: found '%s'\n", __func__, [path_lib_default UTF8String]);
+
+            NSDictionary * atts = [[NSFileManager defaultManager] attributesOfItemAtPath:path_lib_default error:&error];
+            if (atts && atts[NSFileType] == NSFileTypeSymbolicLink) {
+                // Optionally, if this is a symlink, try to resolve it.
+                path_lib_default = [[NSFileManager defaultManager] destinationOfSymbolicLinkAtPath:path_lib_default error:&error];
+                if (path_lib_default && [path_lib_default length] > 0 && ![[path_lib_default substringToIndex:1] isEqualToString:@"/"]) {
+                    // It is a relative path, adding the binary directory as directory prefix.
+                    path_lib_default = [NSString pathWithComponents:@[bin_dir, path_lib_default]];
+                }
+                if (!path_lib_default || ![[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
+                    // Link to the resource could not be resolved.
+                    path_lib_default = nil;
+                } else {
+                    GGML_LOG_INFO("%s: symlink resolved '%s'\n", __func__, [path_lib_default UTF8String]);
+                }
+            }
+        } else {
+            // The resource couldn't be found in the binary's directory.
+            path_lib_default = nil;
+        }
+
+        path_lib = path_lib_default;
+    }
+
+    if (path_lib != nil) {
+        // pre-compiled library found: a single combined default.metallib
+        NSURL * libURL = [NSURL fileURLWithPath:path_lib];
+        GGML_LOG_INFO("%s: loading '%s'\n", __func__, [path_lib UTF8String]);
+
+        res->objs[0]        = [device newLibraryWithURL:libURL error:&error];
+        res->single_library = true;
+        if (!res->objs[0]) {
+            GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
+            ggml_metal_library_free(res);
+            return NULL;
+        }
+
+        GGML_LOG_INFO("%s: loaded in %.3f sec\n", __func__, (ggml_time_us() - t_start) / 1e6);
+        return res;
+    }
+
+    // no pre-compiled metallib: fall back to compiling each kernel source separately
+    GGML_LOG_INFO("%s: default.metallib not found, loading kernel sources\n", __func__);
+
+    NSString * path_resource = [[NSProcessInfo processInfo].environment objectForKey:@"GGML_METAL_PATH_RESOURCES"];
+    if (path_resource) {
+        GGML_LOG_INFO("%s: GGML_METAL_PATH_RESOURCES = %s\n", __func__, [path_resource UTF8String]);
+    }
+
+    // resolve each kind's source path up front (file lookup/logging stays on the calling thread)
+    NSString ** path_per_kind = calloc(GGML_METAL_LIB_COUNT, sizeof(NSString *));
+    for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
+        NSString * rel = [NSString stringWithFormat:@"kernels/%s.metal", k_lib_names[kind]];
+
+        NSString * path_source = nil;
+        if (path_resource) {
+            path_source = [path_resource stringByAppendingPathComponent:rel];
+        } else {
+            NSString * stem = [NSString stringWithFormat:@"kernels/%s", k_lib_names[kind]];
+            path_source = [bundle pathForResource:stem ofType:@"metal"];
+        }
+
+        if (path_source == nil || ![[NSFileManager defaultManager] isReadableFileAtPath:path_source]) {
+            GGML_LOG_WARN("%s: could not locate %s in bundle, falling back to cwd\n", __func__, [rel UTF8String]);
+            path_source = rel;
+        }
+
+        GGML_LOG_DEBUG("%s: loading '%s'\n", __func__, [path_source UTF8String]);
+
+        path_per_kind[kind] = [path_source retain];
+    }
+
+    const bool ok = ggml_metal_library_compile_all(res, device, prep,
+        ^NSString * (int kind, NSError ** err) {
+            return ggml_metal_library_flatten_source(path_per_kind[kind], err);
+        }, "source");
+
+    for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
+        [path_per_kind[kind] release];
+    }
+    free(path_per_kind);
+
+    if (!ok) {
+        ggml_metal_library_free(res);
+        return NULL;
+    }
+
+    return res;
+#endif
 }
 
 ggml_metal_library_t ggml_metal_library_init_from_source(ggml_metal_device_t dev, const char * source, bool verbose) {
@@ -318,10 +589,11 @@ ggml_metal_library_t ggml_metal_library_init_from_source(ggml_metal_device_t dev
         return NULL;
     }
 
-    res->obj       = library;
-    res->dev       = dev;
-    res->pipelines = ggml_metal_pipelines_init();
-    res->lock      = [NSLock new];
+    res->objs[0]        = library;
+    res->single_library = true;
+    res->dev            = dev;
+    res->pipelines      = ggml_metal_pipelines_init();
+    res->lock           = [NSLock new];
 
     return res;
 }
@@ -331,8 +603,14 @@ void ggml_metal_library_free(ggml_metal_library_t lib) {
         return;
     }
 
-    if (lib->obj) {
-        [lib->obj release];
+    for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
+        if (lib->objs[kind]) {
+            [lib->objs[kind] release];
+        }
+    }
+
+    if (lib->fn_to_lib) {
+        [lib->fn_to_lib release];
     }
 
     ggml_metal_pipelines_free(lib->pipelines);
@@ -393,11 +671,28 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_compile_pipeline(ggml_
 
         GGML_LOG_DEBUG("%s: compiling pipeline: base = '%s', name = '%s'\n", __func__, base, name);
 
+        // route to the library that actually defines this kernel; fn_to_lib is
+        // built from -[MTLLibrary functionNames] so it's always in sync
+        int lib_idx = 0;
+        if (!lib->single_library) {
+            NSNumber * idx = lib->fn_to_lib[base_func];
+            if (!idx) {
+                [lib->lock unlock];
+
+                GGML_LOG_ERROR("%s: kernel not found in any metal library: base = '%s', name = '%s'\n", __func__, base, name);
+
+                return res;
+            }
+            lib_idx = [idx intValue];
+        }
+
+        id<MTLLibrary> mtl_lib = lib->objs[lib_idx];
+
         id<MTLFunction> mtl_function;
         if (!cv) {
-            mtl_function = [lib->obj newFunctionWithName:base_func];
+            mtl_function = [mtl_lib newFunctionWithName:base_func];
         } else {
-            mtl_function = [lib->obj newFunctionWithName:base_func constantValues:cv->obj error:&error];
+            mtl_function = [mtl_lib newFunctionWithName:base_func constantValues:cv->obj error:&error];
         }
         if (!mtl_function) {
             [lib->lock unlock];

ggml/src/ggml-metal/kernels/argsort.metal

@@ -0,0 +1,232 @@
+#include "common.h"
+
+// bitonic sort implementation following the CUDA kernels as reference
+typedef void (argsort_t)(
+        constant   ggml_metal_kargs_argsort & args,
+        device   const char * src0,
+        device      int32_t * dst,
+        threadgroup int32_t * shmem_i32 [[threadgroup(0)]],
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]);
+
+template<ggml_sort_order order>
+kernel void kernel_argsort_f32_i32(
+        constant   ggml_metal_kargs_argsort & args,
+        device   const char * src0,
+        device      int32_t * dst,
+        threadgroup int32_t * shmem_i32 [[threadgroup(0)]],
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    // bitonic sort
+    const int col = tpitg[0];
+    const int ib  = tgpig[0] / args.ne01;
+
+    const int i00 = ib*ntg.x;
+    const int i01 = tgpig[0] % args.ne01;
+    const int i02 = tgpig[1];
+    const int i03 = tgpig[2];
+
+    device const float * src0_row = (device const float *) (src0 + args.nb01*i01 + args.nb02*i02 + args.nb03*i03);
+
+    // initialize indices
+    shmem_i32[col] = i00 + col;
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    for (int k = 2; k <= ntg.x; k *= 2) {
+        for (int j = k / 2; j > 0; j /= 2) {
+            int ixj = col ^ j;
+            if (ixj > col) {
+                if ((col & k) == 0) {
+                    if (shmem_i32[col] >= args.ne00 ||
+                       (shmem_i32[ixj] <  args.ne00 && (order == GGML_SORT_ORDER_ASC ?
+                            src0_row[shmem_i32[col]] > src0_row[shmem_i32[ixj]] :
+                            src0_row[shmem_i32[col]] < src0_row[shmem_i32[ixj]]))
+                    ) {
+                        SWAP(shmem_i32[col], shmem_i32[ixj]);
+                    }
+                } else {
+                    if (shmem_i32[ixj] >= args.ne00 ||
+                       (shmem_i32[col] <  args.ne00 && (order == GGML_SORT_ORDER_ASC ?
+                            src0_row[shmem_i32[col]] < src0_row[shmem_i32[ixj]] :
+                            src0_row[shmem_i32[col]] > src0_row[shmem_i32[ixj]]))
+                    ) {
+                        SWAP(shmem_i32[col], shmem_i32[ixj]);
+                    }
+                }
+            }
+
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+        }
+    }
+
+    const int64_t i0 = ib*args.top_k;
+
+    // copy the result to dst without the padding
+    if (i0 + col < args.ne0 && col < args.top_k) {
+        dst += i0 + args.ne0*i01 + args.ne0*args.ne1*i02 + args.ne0*args.ne1*args.ne2*i03;
+
+        dst[col] = shmem_i32[col];
+    }
+}
+
+template [[host_name("kernel_argsort_f32_i32_asc")]]  kernel argsort_t kernel_argsort_f32_i32<GGML_SORT_ORDER_ASC>;
+template [[host_name("kernel_argsort_f32_i32_desc")]] kernel argsort_t kernel_argsort_f32_i32<GGML_SORT_ORDER_DESC>;
+
+typedef void (argsort_merge_t)(
+        constant   ggml_metal_kargs_argsort_merge & args,
+        device const char    * src0,
+        device const int32_t * tmp,
+        device       int32_t * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]);
+
+template<ggml_sort_order order>
+kernel void kernel_argsort_merge_f32_i32(
+        constant   ggml_metal_kargs_argsort_merge & args,
+        device const char    * src0,
+        device const int32_t * tmp,
+        device       int32_t * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+
+    const int im  = tgpig[0] / args.ne01;
+    const int i01 = tgpig[0] % args.ne01;
+    const int i02 = tgpig[1];
+    const int i03 = tgpig[2];
+
+    const int start = im * (2 * args.len);
+
+    const int len0 = MIN(args.len, MAX(0, args.ne0 - (int)(start)));
+    const int len1 = MIN(args.len, MAX(0, args.ne0 - (int)(start + args.len)));
+
+    const int total = len0 + len1;
+
+    device const int32_t * tmp0 = tmp + start
+        + i01*args.ne0
+        + i02*args.ne0*args.ne01
+        + i03*args.ne0*args.ne01*args.ne02;
+
+    device const int32_t * tmp1 = tmp0 + args.len;
+
+    dst += start
+        + i01*args.top_k
+        + i02*args.top_k*args.ne01
+        + i03*args.top_k*args.ne01*args.ne02;
+
+    device const float * src0_row = (device const float *)(src0
+        + args.nb01*i01
+        + args.nb02*i02
+        + args.nb03*i03);
+
+    if (total == 0) {
+        return;
+    }
+
+    const int chunk = (total + ntg.x - 1) / ntg.x;
+
+    const int k0 = tpitg.x * chunk;
+    const int k1 = MIN(MIN(k0 + chunk, total), args.top_k);
+
+    if (k0 >= args.top_k) {
+        return;
+    }
+
+    if (k0 >= total) {
+        return;
+    }
+
+    int low  = k0 > len1 ? k0 - len1 : 0;
+    int high = MIN(k0, len0);
+
+    // binary-search partition (i, j) such that i + j = k
+    while (low < high) {
+        const int mid = (low + high) >> 1;
+
+        const int32_t idx0 = tmp0[mid];
+        const int32_t idx1 = tmp1[k0 - mid - 1];
+
+        const float val0 = src0_row[idx0];
+        const float val1 = src0_row[idx1];
+
+        bool take_left;
+        if (order == GGML_SORT_ORDER_ASC) {
+            take_left = (val0 <= val1);
+        } else {
+            take_left = (val0 >= val1);
+        }
+
+        if (take_left) {
+            low = mid + 1;
+        } else {
+            high = mid;
+        }
+    }
+
+    int i = low;
+    int j = k0 - i;
+
+    // keep the merge fronts into registers
+    int32_t idx0 = 0;
+    float   val0 = 0.0f;
+    if (i < len0) {
+        idx0 = tmp0[i];
+        val0 = src0_row[idx0];
+    }
+
+    int32_t idx1 = 0;
+    float   val1 = 0.0f;
+    if (j < len1) {
+        idx1 = tmp1[j];
+        val1 = src0_row[idx1];
+    }
+
+    for (int k = k0; k < k1; ++k) {
+        int32_t out_idx;
+
+        if (i >= len0) {
+            while (k < k1) {
+                dst[k++] = tmp1[j++];
+            }
+            break;
+        } else if (j >= len1) {
+            while (k < k1) {
+                dst[k++] = tmp0[i++];
+            }
+            break;
+        } else {
+            bool take_left;
+
+            if (order == GGML_SORT_ORDER_ASC) {
+                take_left = (val0 <= val1);
+            } else {
+                take_left = (val0 >= val1);
+            }
+
+            if (take_left) {
+                out_idx = idx0;
+                ++i;
+                if (i < len0) {
+                    idx0 = tmp0[i];
+                    val0 = src0_row[idx0];
+                }
+            } else {
+                out_idx = idx1;
+                ++j;
+                if (j < len1) {
+                    idx1 = tmp1[j];
+                    val1 = src0_row[idx1];
+                }
+            }
+        }
+
+        dst[k] = out_idx;
+    }
+}
+
+template [[host_name("kernel_argsort_merge_f32_i32_asc")]]  kernel argsort_merge_t kernel_argsort_merge_f32_i32<GGML_SORT_ORDER_ASC>;
+template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<GGML_SORT_ORDER_DESC>;

ggml/src/ggml-metal/kernels/binbcast.metal

@@ -0,0 +1,226 @@
+#include "common.h"
+
+// OP: 0 - add, 1 - sub, 2 - mul, 3 - div
+constant short FC_bin_op [[function_constant(FC_BIN + 0)]];
+constant short FC_bin_f  [[function_constant(FC_BIN + 1)]];
+constant bool  FC_bin_rb [[function_constant(FC_BIN + 2)]];
+constant bool  FC_bin_cb [[function_constant(FC_BIN + 3)]];
+
+template <typename T0, typename T1, typename T>
+kernel void kernel_bin_fuse_impl(
+        constant ggml_metal_kargs_bin & args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+#define FC_OP FC_bin_op
+#define FC_F  FC_bin_f
+#define FC_RB FC_bin_rb
+#define FC_CB FC_bin_cb
+
+    if (FC_RB) {
+        // row broadcast
+        const uint i0 = tgpig.y*args.ne00 + tgpig.x;
+        const uint i1 = FC_CB ? tgpig.x%args.ne10 : tgpig.x;
+
+        device const T0 * src0_row = (device const T0 *) (src0);
+        device       T  * dst_row  = (device       T  *) (dst);
+
+        if (FC_F == 1) {
+            device const T1 * src1_row = (device const T1 *) (src1 + args.o1[0]);
+
+            if (FC_OP == 0) {
+                dst_row[i0] = src0_row[i0] + src1_row[i1];
+            }
+
+            if (FC_OP == 1) {
+                dst_row[i0] = src0_row[i0] - src1_row[i1];
+            }
+
+            if (FC_OP == 2) {
+                dst_row[i0] = src0_row[i0] * src1_row[i1];
+            }
+
+            if (FC_OP == 3) {
+                dst_row[i0] = src0_row[i0] / src1_row[i1];
+            }
+        } else {
+            T0 res = src0_row[i0];
+
+            if (FC_OP == 0) {
+                FOR_UNROLL (short j = 0; j < FC_F; ++j) {
+                    res += ((device const T1 *) (src1 + args.o1[j]))[i1];
+                }
+            }
+
+            if (FC_OP == 1) {
+                FOR_UNROLL (short j = 0; j < FC_F; ++j) {
+                    res -= ((device const T1 *) (src1 + args.o1[j]))[i1];
+                }
+            }
+
+            if (FC_OP == 2) {
+                FOR_UNROLL (short j = 0; j < FC_F; ++j) {
+                    res *= ((device const T1 *) (src1 + args.o1[j]))[i1];
+                }
+            }
+
+            if (FC_OP == 3) {
+                FOR_UNROLL (short j = 0; j < FC_F; ++j) {
+                    res /= ((device const T1 *) (src1 + args.o1[j]))[i1];
+                }
+            }
+
+            dst_row[i0] = res;
+        }
+    } else {
+        const int i03 = tgpig.z;
+        const int i02 = tgpig.y;
+        const int i01 = tgpig.x;
+
+        if (i01 >= args.ne01) {
+            return;
+        }
+
+        const int i13 = i03%args.ne13;
+        const int i12 = i02%args.ne12;
+        const int i11 = i01%args.ne11;
+
+        device const T0 * src0_ptr = (device const T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + args.offs);
+        device       T  * dst_ptr  = (device       T  *) (dst  + i03*args.nb3  + i02*args.nb2  + i01*args.nb1  + args.offs);
+
+        if (FC_F == 1) {
+            device const T1 * src1_ptr = (device const T1 *) (src1 + args.o1[0] + i13*args.nb13 + i12*args.nb12 + i11*args.nb11);
+
+            for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
+                const int i10 = FC_CB ? i0%args.ne10 : i0;
+
+                if (FC_OP == 0) {
+                    dst_ptr[i0] = src0_ptr[i0] + src1_ptr[i10];
+                }
+
+                if (FC_OP == 1) {
+                    dst_ptr[i0] = src0_ptr[i0] - src1_ptr[i10];
+                }
+
+                if (FC_OP == 2) {
+                    dst_ptr[i0] = src0_ptr[i0] * src1_ptr[i10];
+                }
+
+                if (FC_OP == 3) {
+                    dst_ptr[i0] = src0_ptr[i0] / src1_ptr[i10];
+                }
+            }
+        } else {
+            device const T1 * src1_ptr[8];
+            FOR_UNROLL (short j = 0; j < FC_F; ++j) {
+                src1_ptr[j] = (device const T1 *) (src1 + args.o1[j] + i13*args.nb13 + i12*args.nb12 + i11*args.nb11);
+            }
+
+            for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
+                const int i10 = FC_CB ? i0%args.ne10 : i0;
+
+                T res = src0_ptr[i0];
+
+                if (FC_OP == 0) {
+                    FOR_UNROLL (short j = 0; j < FC_F; ++j) {
+                        res += src1_ptr[j][i10];
+                    }
+                }
+
+                if (FC_OP == 1) {
+                    FOR_UNROLL (short j = 0; j < FC_F; ++j) {
+                        res -= src1_ptr[j][i10];
+                    }
+                }
+
+                if (FC_OP == 2) {
+                    FOR_UNROLL (short j = 0; j < FC_F; ++j) {
+                        res *= src1_ptr[j][i10];
+                    }
+                }
+
+                if (FC_OP == 3) {
+                    FOR_UNROLL (short j = 0; j < FC_F; ++j) {
+                        res /= src1_ptr[j][i10];
+                    }
+                }
+
+                dst_ptr[i0] = res;
+            }
+        }
+    }
+
+#undef FC_OP
+#undef FC_F
+#undef FC_RB
+#undef FC_CB
+}
+
+typedef decltype(kernel_bin_fuse_impl<float, float, float>) kernel_bin_fuse_t;
+
+template [[host_name("kernel_bin_fuse_f32_f32_f32")]]   kernel kernel_bin_fuse_t kernel_bin_fuse_impl<float,  float,  float>;
+template [[host_name("kernel_bin_fuse_f32_f32_f32_4")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl<float4, float4, float4>;
+
+kernel void kernel_add_id(
+        constant ggml_metal_kargs_add_id & args,
+        device const char * src0,
+        device const char * src1,
+        device const char * src2,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    const int i1 = tgpig.x;
+    const int i2 = tgpig.y;
+
+    const int i11 = *((device const int32_t *) (src2 + i1*sizeof(int32_t) + i2*args.nb21));
+
+    const size_t nb1 = args.ne0 * sizeof(float);
+    const size_t nb2 = args.ne1 * nb1;
+
+    device       float * dst_row  = (device       float *)((device char *)dst  +  i1*nb1       + i2*nb2);
+    device const float * src0_row = (device const float *)((device char *)src0 +  i1*args.nb01 + i2*args.nb02);
+    device const float * src1_row = (device const float *)((device char *)src1 + i11*args.nb11);
+
+    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
+        dst_row[i0] = src0_row[i0] + src1_row[i0];
+    }
+}
+
+template<typename T>
+kernel void kernel_repeat(
+        constant ggml_metal_kargs_repeat & args,
+        device const char * src0,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    const int i3 = tgpig.z;
+    const int i2 = tgpig.y;
+    const int i1 = tgpig.x;
+
+    const int i03 = i3%args.ne03;
+    const int i02 = i2%args.ne02;
+    const int i01 = i1%args.ne01;
+
+    device const char * src0_ptr = src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01;
+    device       char * dst_ptr  = dst  +  i3*args.nb3  +  i2*args.nb2  +  i1*args.nb1;
+
+    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
+        const int i00 = i0%args.ne00;
+        *((device T *)(dst_ptr + i0*args.nb0)) = *((device T *)(src0_ptr + i00*args.nb00));
+    }
+}
+
+typedef decltype(kernel_repeat<float>) kernel_repeat_t;
+
+template [[host_name("kernel_repeat_f32")]] kernel kernel_repeat_t kernel_repeat<float>;
+template [[host_name("kernel_repeat_f16")]] kernel kernel_repeat_t kernel_repeat<half>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_repeat_bf16")]] kernel kernel_repeat_t kernel_repeat<bfloat>;
+#endif
+template [[host_name("kernel_repeat_i32")]] kernel kernel_repeat_t kernel_repeat<int>;
+template [[host_name("kernel_repeat_i16")]] kernel kernel_repeat_t kernel_repeat<short>;

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

@@ -0,0 +1,126 @@
+#pragma once
+
+#include "ggml-metal-impl.h"
+
+#include <metal_stdlib>
+
+#ifdef GGML_METAL_HAS_TENSOR
+#include <metal_tensor>
+
+#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
+#endif
+
+using namespace metal;
+
+#define MAX(x, y) ((x) > (y) ? (x) : (y))
+#define MIN(x, y) ((x) < (y) ? (x) : (y))
+#define SWAP(x, y) { auto tmp = (x); (x) = (y); (y) = tmp; }
+
+#define PAD2(x, n) (((x) + (n) - 1) & ~((n) - 1))
+
+#define FOR_UNROLL(x) _Pragma("clang loop unroll(full)") for (x)
+
+#define N_SIMDWIDTH 32 // assuming SIMD group size is 32
+
+// ref: https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf
+//
+// cmd:
+//   .../usr/bin/metal -dM -E -c                             ggml/src/ggml-metal/kernels/<src>.metal
+//   .../usr/bin/metal -dM -E -c -target air64-apple-ios14.0 ggml/src/ggml-metal/kernels/<src>.metal
+//
+#if __METAL_VERSION__ < 310 && defined(GGML_METAL_HAS_BF16)
+#undef GGML_METAL_HAS_BF16
+#endif
+
+#if defined(GGML_METAL_HAS_BF16)
+typedef matrix<bfloat, 4, 4> bfloat4x4;
+typedef matrix<bfloat, 2, 4> bfloat2x4;
+#endif
+
+constexpr constant static float kvalues_iq4nl_f[16] = {
+    -127.f, -104.f, -83.f, -65.f, -49.f, -35.f, -22.f, -10.f, 1.f, 13.f, 25.f, 38.f, 53.f, 69.f, 89.f, 113.f
+};
+
+constexpr constant static float kvalues_mxfp4_f[16] = {
+    0, .5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f, -0, -.5f, -1.f, -1.5f, -2.f, -3.f, -4.f, -6.f
+};
+
+static inline int best_index_int8(int n, constant float * val, float x) {
+    if (x <= val[0]) return 0;
+    if (x >= val[n-1]) return n-1;
+    int ml = 0, mu = n-1;
+    while (mu-ml > 1) {
+        int mav = (ml+mu)/2;
+        if (x < val[mav]) mu = mav; else ml = mav;
+    }
+    return x - val[mu-1] < val[mu] - x ? mu-1 : mu;
+}
+
+static inline float e8m0_to_fp32(uint8_t x) {
+    uint32_t bits;
+
+    if (x == 0) {
+        bits = 0x00400000;
+    } else {
+        bits = (uint32_t) x << 23;
+    }
+
+    return as_type<float>(bits);
+}
+
+static inline float dot(float x, float y) {
+    return x*y;
+}
+
+static inline float sum(float x) {
+    return x;
+}
+
+static inline float sum(float4 x) {
+    return x[0] + x[1] + x[2] + x[3];
+}
+
+enum ggml_sort_order {
+    GGML_SORT_ORDER_ASC,
+    GGML_SORT_ORDER_DESC,
+};
+
+constant float GELU_COEF_A     = 0.044715f;
+constant float GELU_QUICK_COEF = -1.702f;
+constant float SQRT_2_OVER_PI  = 0.79788456080286535587989211986876f;
+constant float SQRT_2_INV      = 0.70710678118654752440084436210484f;
+
+// based on Abramowitz and Stegun formula 7.1.26 or similar Hastings' approximation
+// ref: https://www.johndcook.com/blog/python_erf/
+constant float p_erf  = 0.3275911f;
+constant float a1_erf = 0.254829592f;
+constant float a2_erf = -0.284496736f;
+constant float a3_erf = 1.421413741f;
+constant float a4_erf = -1.453152027f;
+constant float a5_erf = 1.061405429f;
+
+template<typename T>
+inline T erf_approx(T x) {
+    T sign_x = sign(x);
+    x = fabs(x);
+    T t = 1.0f / (1.0f + p_erf * x);
+    T y = 1.0f - (((((a5_erf * t + a4_erf) * t) + a3_erf) * t + a2_erf) * t + a1_erf) * t * exp(-x * x);
+    return sign_x * y;
+}
+
+template<typename T> T elu_approx(T x);
+
+template<> inline float elu_approx<float>(float x) {
+    return (x > 0.f) ? x : (exp(x) - 1);
+}
+
+template<> inline float4 elu_approx<float4>(float4 x) {
+    float4 res;
+
+    res[0] = (x[0] > 0.0f) ? x[0] : (exp(x[0]) - 1.0f);
+    res[1] = (x[1] > 0.0f) ? x[1] : (exp(x[1]) - 1.0f);
+    res[2] = (x[2] > 0.0f) ? x[2] : (exp(x[2]) - 1.0f);
+    res[3] = (x[3] > 0.0f) ? x[3] : (exp(x[3]) - 1.0f);
+
+    return res;
+}

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

@@ -0,0 +1,723 @@
+#include "common.h"
+
+typedef void (im2col_t)(
+        constant ggml_metal_kargs_im2col & args,
+        device const float * x,
+        device        char * dst,
+        uint3 tgpig[[threadgroup_position_in_grid]],
+        uint3  tgpg[[threadgroups_per_grid]],
+        uint3 tpitg[[thread_position_in_threadgroup]],
+        uint3   ntg[[threads_per_threadgroup]]);
+
+template <typename T>
+kernel void kernel_im2col(
+        constant ggml_metal_kargs_im2col & args,
+        device const float * x,
+        device        char * dst,
+        uint3 tgpig[[threadgroup_position_in_grid]],
+        uint3  tgpg[[threadgroups_per_grid]],
+        uint3 tpitg[[thread_position_in_threadgroup]],
+        uint3   ntg[[threads_per_threadgroup]]) {
+//    const int64_t IC = tgpg[0];
+    const int64_t OH = tgpg[1];
+    const int64_t OW = tgpg[2];
+
+    const int64_t KH = ntg[1];
+    const int64_t KW = ntg[2];
+
+          int64_t in  = tpitg[0];
+    const int64_t ikh = tpitg[1];
+    const int64_t ikw = tpitg[2];
+
+    const int64_t iic = tgpig[0];
+    const int64_t ioh = tgpig[1];
+    const int64_t iow = tgpig[2];
+
+    const int64_t iiw = iow*args.s0 + ikw*args.d0 - args.p0;
+    const int64_t iih = ioh*args.s1 + ikh*args.d1 - args.p1;
+
+    int64_t offset_dst = (in*OH*OW + ioh*OW + iow)*args.CHW + (iic*(KH*KW) + ikh*KW + ikw);
+
+    device T * pdst = (device T *) (dst);
+
+    if (iih < 0 || iih >= args.IH || iiw < 0 || iiw >= args.IW) {
+        while (in < args.N) {
+            pdst[offset_dst] = 0.0f;
+            offset_dst += ntg[0]*args.CHW*OH*OW;
+
+            in += ntg[0];
+        }
+    } else {
+        int64_t offset_src = in*args.ofs0 + iic*args.ofs1 + iih*args.IW + iiw;
+
+        while (in < args.N) {
+            pdst[offset_dst] = x[offset_src];
+
+            offset_dst += ntg[0]*args.CHW*OH*OW;
+            offset_src += ntg[0]*args.ofs0;
+
+            in += ntg[0];
+        }
+    }
+}
+
+template [[host_name("kernel_im2col_f32")]] kernel im2col_t kernel_im2col<float>;
+template [[host_name("kernel_im2col_f16")]] kernel im2col_t kernel_im2col<half>;
+
+// TODO: optimize
+typedef void (im2col_ext_t)(
+        constant ggml_metal_kargs_im2col & args,
+        device const float * x,
+        device        char * dst,
+        uint3 tgpig[[threadgroup_position_in_grid]],
+        uint3  tgpg[[threadgroups_per_grid]],
+        uint3 tpitg[[thread_position_in_threadgroup]],
+        uint3   ntg[[threads_per_threadgroup]]);
+
+template <typename T>
+kernel void kernel_im2col_ext(
+        constant ggml_metal_kargs_im2col & args,
+        device const float * x,
+        device        char * dst,
+        uint3 tgpig[[threadgroup_position_in_grid]],
+        uint3  tgpg[[threadgroups_per_grid]],      // tgpg[0] = D x IC x KH x KW, CHW = IC x KH x KW
+        uint3 tpitg[[thread_position_in_threadgroup]],
+        uint3   ntg[[threads_per_threadgroup]]) {  // [M, 1, 1]
+    const int64_t KHW = (int64_t)args.KHW;
+
+    const int64_t d   = tgpig[0] / args.CHW;
+    const int64_t chw = tgpig[0] % args.CHW;
+    const int64_t tgpig_0 = chw / KHW;  // 0 ~ (IC - 1)
+    const int64_t HW = tgpig[0] % KHW;
+
+    const int64_t tpitg_0 = (d * ntg[0]) + tpitg[0];
+    if (tpitg_0 >= args.N) {
+        return;
+    }
+
+    const int64_t tpitg_1 = HW / args.KW;
+    const int64_t tpitg_2 = HW % args.KW;
+
+    const int64_t iiw = tgpig[2] * args.s0 + tpitg_2 * args.d0 - args.p0;
+    const int64_t iih = tgpig[1] * args.s1 + tpitg_1 * args.d1 - args.p1;
+
+    const int64_t offset_dst =
+        (tpitg_0 * tgpg[1] * tgpg[2] + tgpig[1] * tgpg[2] + tgpig[2]) * args.CHW +
+        (tgpig_0 * KHW + tpitg_1 * args.KW + tpitg_2);
+
+    device T * pdst = (device T *) (dst);
+
+    if (iih < 0 || iih >= args.IH || iiw < 0 || iiw >= args.IW) {
+        pdst[offset_dst] = 0.0f;
+    } else {
+        const int64_t offset_src = tpitg_0 * args.ofs0 + tgpig_0 * args.ofs1;
+        pdst[offset_dst] = x[offset_src + iih * args.IW + iiw];
+    }
+}
+
+template [[host_name("kernel_im2col_ext_f32")]] kernel im2col_ext_t kernel_im2col_ext<float>;
+template [[host_name("kernel_im2col_ext_f16")]] kernel im2col_ext_t kernel_im2col_ext<half>;
+
+template <typename T>
+kernel void kernel_col2im_1d(
+        constant ggml_metal_kargs_col2im_1d & args,
+        device const T * col,
+        device       T * dst,
+        uint         tgpig [[threadgroup_position_in_grid]],
+        uint         tpitg [[thread_position_in_threadgroup]],
+        uint         ntg   [[threads_per_threadgroup]]) {
+
+    const int idx = tgpig * ntg + tpitg;
+    if (idx >= args.T_out * args.OC) {
+        return;
+    }
+
+    const int t_out = idx % args.T_out;
+    const int oc    = idx / args.T_out;
+    const int t_abs = t_out + args.p0;  // absolute position in uncropped signal
+
+    int t_in_min = (t_abs - args.K + args.s0) / args.s0;  // ceil((t_abs - K + 1) / s0)
+    if (t_in_min < 0) {
+        t_in_min = 0;
+    }
+    int t_in_max = t_abs / args.s0;
+    if (t_in_max >= args.T_in) {
+        t_in_max = args.T_in - 1;
+    }
+
+    float sum = 0.0f;
+    for (int t_in = t_in_min; t_in <= t_in_max; t_in++) {
+        const int k = t_abs - t_in * args.s0;
+        sum += float(col[(oc * args.K + k) + t_in * args.K_OC]);
+    }
+
+    dst[t_out + oc * args.T_out] = T(sum);
+}
+
+template [[host_name("kernel_col2im_1d_f32")]]  kernel void kernel_col2im_1d<float>(constant ggml_metal_kargs_col2im_1d &, device const float *, device float *, uint, uint, uint);
+template [[host_name("kernel_col2im_1d_f16")]]  kernel void kernel_col2im_1d<half>(constant ggml_metal_kargs_col2im_1d &, device const half *, device half *, uint, uint, uint);
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_col2im_1d_bf16")]] kernel void kernel_col2im_1d<bfloat>(constant ggml_metal_kargs_col2im_1d &, device const bfloat *, device bfloat *, uint, uint, uint);
+#endif
+
+template <typename TK>
+kernel void kernel_conv_2d(
+        constant ggml_metal_kargs_conv_2d & args,
+        device const char * weights,
+        device const char * src,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        uint3    tgpg[[threadgroups_per_grid]],
+        uint3   tpitg[[thread_position_in_threadgroup]],
+        uint3     ntg[[threads_per_threadgroup]]) {
+
+    const uint threads_per_tg = ntg.x * ntg.y * ntg.z;
+    const uint tg_index = (tgpig.z * tgpg.y + tgpig.y) * tgpg.x + tgpig.x;
+    const uint local_thread = tpitg.z * (ntg.x * ntg.y) + tpitg.y * ntg.x + tpitg.x;
+    const uint thread_index = tg_index * threads_per_tg + local_thread;
+    const uint64_t total_threads = (uint64_t) threads_per_tg * tgpg.x * tgpg.y * tgpg.z;
+    const uint64_t total_outputs = (uint64_t) args.N * args.OC * args.OH * args.OW;
+
+    for (uint64_t index = thread_index; index < total_outputs; index += total_threads) {
+        uint64_t tmp = index;
+
+        const int32_t ow = tmp % args.OW; tmp /= args.OW;
+        const int32_t oh = tmp % args.OH; tmp /= args.OH;
+        const int32_t oc = tmp % args.OC; tmp /= args.OC;
+        const int32_t  n = tmp;
+
+        float acc = 0.0f;
+
+        const int32_t base_x = ow*args.s0 - args.p0;
+        const int32_t base_y = oh*args.s1 - args.p1;
+
+        int32_t ky_start = 0;
+        if (base_y < 0) {
+            ky_start = (-base_y + args.d1 - 1)/args.d1;
+        }
+        int32_t ky_end = args.KH;
+        const int32_t y_max = args.IH - 1 - base_y;
+        if (y_max < 0) {
+            ky_end = ky_start;
+        } else if (base_y + (args.KH - 1)*args.d1 >= args.IH) {
+            ky_end = min(ky_end, y_max/args.d1 + 1);
+        }
+
+        int32_t kx_start = 0;
+        if (base_x < 0) {
+            kx_start = (-base_x + args.d0 - 1)/args.d0;
+        }
+        int32_t kx_end = args.KW;
+        const int32_t x_max = args.IW - 1 - base_x;
+        if (x_max < 0) {
+            kx_end = kx_start;
+        } else if (base_x + (args.KW - 1)*args.d0 >= args.IW) {
+            kx_end = min(kx_end, x_max/args.d0 + 1);
+        }
+
+        if (ky_start < ky_end && kx_start < kx_end) {
+            const uint64_t src_base_n = (uint64_t) n  * args.nb13;
+            const uint64_t w_base_oc  = (uint64_t) oc * args.nb03;
+
+            for (int32_t ic = 0; ic < args.IC; ++ic) {
+                const uint64_t src_base_nc = src_base_n + (uint64_t) ic * args.nb12;
+                const uint64_t w_base_ocic = w_base_oc  + (uint64_t) ic * args.nb02;
+
+                for (int32_t ky = ky_start; ky < ky_end; ++ky) {
+                    const int32_t iy = base_y + ky*args.d1;
+                    const uint64_t src_base_row = src_base_nc + (uint64_t) iy * args.nb11;
+                    const uint64_t w_base_row   = w_base_ocic + (uint64_t) ky * args.nb01;
+
+                    for (int32_t kx = kx_start; kx < kx_end; ++kx) {
+                        const int32_t ix = base_x + kx*args.d0;
+                        const uint64_t src_offs = src_base_row + (uint64_t) ix * args.nb10;
+                        const uint64_t w_offs   = w_base_row   + (uint64_t) kx * args.nb00;
+
+                        const float x = *(device const float *)(src + src_offs);
+                        const float w = (float) (*(device const TK *)(weights + w_offs));
+
+                        acc += x * w;
+                    }
+                }
+            }
+        }
+
+        const uint64_t dst_offs =
+            (uint64_t) n  * args.nb3 +
+            (uint64_t) oc * args.nb2 +
+            (uint64_t) oh * args.nb1 +
+            (uint64_t) ow * args.nb0;
+
+        *(device float *)(dst + dst_offs) = acc;
+    }
+}
+
+template [[host_name("kernel_conv_2d_f32_f32")]]
+kernel void kernel_conv_2d<float>(
+        constant ggml_metal_kargs_conv_2d & args,
+        device const char * weights,
+        device const char * src,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        uint3    tgpg[[threadgroups_per_grid]],
+        uint3   tpitg[[thread_position_in_threadgroup]],
+        uint3     ntg[[threads_per_threadgroup]]);
+
+template [[host_name("kernel_conv_2d_f16_f32")]]
+kernel void kernel_conv_2d<half>(
+        constant ggml_metal_kargs_conv_2d & args,
+        device const char * weights,
+        device const char * src,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        uint3    tgpg[[threadgroups_per_grid]],
+        uint3   tpitg[[thread_position_in_threadgroup]],
+        uint3     ntg[[threads_per_threadgroup]]);
+
+typedef void (conv_transpose_1d_t)(
+        constant ggml_metal_kargs_conv_transpose_1d & args,
+        device const float * src0,
+        device const float * src1,
+        device        char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        uint3    tgpg[[threadgroups_per_grid]]);
+
+template <typename T>
+kernel void kernel_conv_transpose_1d(
+        constant ggml_metal_kargs_conv_transpose_1d & args,
+        device const     T * src0,
+        device const float * src1,
+        device        char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        uint3   tgpg[[threadgroups_per_grid]]) {
+
+    // For output position j on the time axis, only input positions
+    //   i such that i*s0 <= j < i*s0 + K
+    // contribute -- i.e. i in [ceil((j - K + 1)/s0), floor(j/s0)]
+    // intersected with [0, IL-1]. That's at most ceil(K/s0) values
+    // (typically 2 for stride==K/2 transposed convs).
+    const int32_t j  = tgpig[0];
+    const int32_t s0 = args.s0;
+    const int32_t K  = args.K;
+    const int32_t IL = args.IL;
+
+    int32_t i_min;
+    {
+        int32_t a = j - K + 1;
+        i_min = a <= 0 ? 0 : (a + s0 - 1) / s0; // ceil(a/s0) for a>0
+    }
+    int32_t i_max = j / s0;
+    if (i_max > IL - 1) i_max = IL - 1;
+
+    float v = 0.0f;
+    if (i_min <= i_max) {
+        for (int64_t c = 0; c < args.IC; c++) {
+            const int32_t kernel_offset = c * tgpg[1] * K + K * tgpig[1];
+            const int32_t input_offset  = c * IL;
+
+            for (int32_t i = i_min; i <= i_max; i++) {
+                v += float(src0[kernel_offset + j - i * s0]) * src1[input_offset + i];
+            }
+        }
+    }
+
+    device float * dst_ptr = (device float *) (dst + tgpig[0] * args.nb0 + tgpig[1] * args.nb1);
+
+    dst_ptr[0] = v;
+}
+
+template [[host_name("kernel_conv_transpose_1d_f32_f32")]]
+kernel void kernel_conv_transpose_1d<float>(
+    constant ggml_metal_kargs_conv_transpose_1d & args,
+    device const float * src0,
+    device const float * src1,
+    device        char * dst,
+    uint3   tgpig[[threadgroup_position_in_grid]],
+    uint3    tgpg[[threadgroups_per_grid]]);
+
+template [[host_name("kernel_conv_transpose_1d_f16_f32")]]
+kernel void kernel_conv_transpose_1d<half>(
+    constant ggml_metal_kargs_conv_transpose_1d & args,
+    device const half  * src0,
+    device const float * src1,
+    device        char * dst,
+    uint3   tgpig[[threadgroup_position_in_grid]],
+    uint3    tgpg[[threadgroups_per_grid]]);
+
+
+typedef void (conv_transpose_2d_t)(
+        constant ggml_metal_kargs_conv_transpose_2d & args,
+        device const float * src0,
+        device const float * src1,
+        device        char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        uint3    tgpg[[threadgroups_per_grid]]);
+
+template <typename T>
+kernel void kernel_conv_transpose_2d(
+        constant ggml_metal_kargs_conv_transpose_2d & args,
+        device const T * src0,
+        device const float * src1,
+        device        char * dst,
+        threadgroup float * shared_sum [[threadgroup(0)]],
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        uint3   tpitg[[thread_position_in_threadgroup]],
+        uint3     ntg[[threads_per_threadgroup]]) {
+
+    const int64_t out_x = tgpig[0];
+    const int64_t out_y = tgpig[1];
+    const int64_t out_c = tgpig[2];
+
+    const int64_t kw = tpitg[0];
+    const int64_t kh = tpitg[1];
+
+    float v = 0.0f;
+
+    for (int64_t in_c = 0; in_c < args.IC; in_c++) {
+        int64_t in_y = out_y - kh;
+
+        if (in_y < 0 || in_y % args.s0) continue;
+
+        in_y /= args.s0;
+
+        if (in_y >= args.IH) continue;
+
+        int64_t in_x = out_x - kw;
+
+        if (in_x < 0 || in_x % args.s0) continue;
+
+        in_x /= args.s0;
+
+        if (in_x >= args.IW) continue;
+
+        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];
+    }
+
+    const uint tid = tpitg.y * ntg.x + tpitg.x;
+    shared_sum[tid] = v;
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    if (tid == 0) {
+        float total = 0.0f;
+        const uint num_threads = ntg.x * ntg.y;
+        for (uint i = 0; i < num_threads; i++) {
+            total += shared_sum[i];
+        }
+
+        device float * dst_ptr = (device float *) (dst + out_x*args.nb0 + out_y * args.nb1 + out_c*args.nb2);
+        dst_ptr[0] = total;
+    }
+}
+
+template [[host_name("kernel_conv_transpose_2d_f32_f32")]]
+kernel void kernel_conv_transpose_2d<float>(
+    constant ggml_metal_kargs_conv_transpose_2d & args,
+    device const float * src0,
+    device const float * src1,
+    device        char * dst,
+    threadgroup float * shared_sum [[threadgroup(0)]],
+    uint3   tgpig[[threadgroup_position_in_grid]],
+    uint3   tpitg[[thread_position_in_threadgroup]],
+    uint3     ntg[[threads_per_threadgroup]]);
+
+template [[host_name("kernel_conv_transpose_2d_f16_f32")]]
+kernel void kernel_conv_transpose_2d<half>(
+    constant ggml_metal_kargs_conv_transpose_2d & args,
+    device const half  * src0,
+    device const float * src1,
+    device        char * dst,
+    threadgroup float * shared_sum [[threadgroup(0)]],
+    uint3   tgpig[[threadgroup_position_in_grid]],
+    uint3   tpitg[[thread_position_in_threadgroup]],
+    uint3     ntg[[threads_per_threadgroup]]);
+
+// grid: x = C tile, y = OH, z = OW * N (for channel-contiguous layouts)
+template <typename TK>
+kernel void kernel_conv_2d_dw_tiled(
+        constant ggml_metal_kargs_conv_2d_dw & args,
+        device const char * weights,
+        device const char * src,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        uint3   tpitg[[thread_position_in_threadgroup]],
+        uint3     ntg[[threads_per_threadgroup]]) {
+
+    const int32_t c = (int32_t)(tgpig.x * ntg.x + tpitg.x);
+    if (c >= args.C) {
+        return;
+    }
+
+    const int32_t oh = tgpig.y;
+    const int32_t own = tgpig.z;
+    const int32_t ow = own % args.OW;
+    const int32_t n  = own / args.OW;
+
+    const int32_t base_y = oh*args.s1 - args.p1;
+
+    int32_t ky_start = 0;
+    if (base_y < 0) {
+        ky_start = (-base_y + args.d1 - 1)/args.d1;
+    }
+    int32_t ky_end = args.KH;
+    const int32_t y_max = args.IH - 1 - base_y;
+    if (y_max < 0) {
+        ky_end = ky_start;
+    } else if (base_y + (args.KH - 1)*args.d1 >= args.IH) {
+        ky_end = min(ky_end, y_max/args.d1 + 1);
+    }
+
+    const int32_t base_x = ow*args.s0 - args.p0;
+
+    int32_t kx_start = 0;
+    if (base_x < 0) {
+        kx_start = (-base_x + args.d0 - 1)/args.d0;
+    }
+    int32_t kx_end = args.KW;
+    const int32_t x_max = args.IW - 1 - base_x;
+    if (x_max < 0) {
+        kx_end = kx_start;
+    } else if (base_x + (args.KW - 1)*args.d0 >= args.IW) {
+        kx_end = min(kx_end, x_max/args.d0 + 1);
+    }
+
+    float acc = 0.0f;
+
+    if (ky_start < ky_end && kx_start < kx_end) {
+        const uint64_t w_base   = (uint64_t) c * args.nb02;
+        const uint64_t src_base = (uint64_t) n * args.nb13 + (uint64_t) c * args.nb12;
+
+        for (int32_t ky = ky_start; ky < ky_end; ++ky) {
+            const int32_t iy = base_y + ky*args.d1;
+            const uint64_t src_row = src_base + (uint64_t) iy * args.nb11;
+            const uint64_t w_row = w_base + (uint64_t) ky * args.nb01;
+
+            for (int32_t kx = kx_start; kx < kx_end; ++kx) {
+                const int32_t ix = base_x + kx*args.d0;
+                const float x = *(device const float *)(src + src_row + (uint64_t) ix * args.nb10);
+                const float w = (float)(*(device const TK *)(weights + w_row + (uint64_t) kx * args.nb00));
+                acc += x * w;
+            }
+        }
+    }
+
+    const uint64_t dst_offs =
+        (uint64_t) n  * args.nb3 +
+        (uint64_t) c  * args.nb2 +
+        (uint64_t) oh * args.nb1 +
+        (uint64_t) ow * args.nb0;
+
+    *(device float *)(dst + dst_offs) = acc;
+}
+
+// grid: x = OW tile, y = OH, z = C * N (for spatially-contiguous layouts)
+template <typename TK>
+kernel void kernel_conv_2d_dw(
+        constant ggml_metal_kargs_conv_2d_dw & args,
+        device const char * weights,
+        device const char * src,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        uint3   tpitg[[thread_position_in_threadgroup]],
+        uint3     ntg[[threads_per_threadgroup]]) {
+
+    const int32_t oh = tgpig.y;
+    const int32_t cn = tgpig.z;
+    const int32_t c  = cn % args.C;
+    const int32_t n  = cn / args.C;
+
+    const int32_t base_y = oh*args.s1 - args.p1;
+
+    int32_t ky_start = 0;
+    if (base_y < 0) {
+        ky_start = (-base_y + args.d1 - 1)/args.d1;
+    }
+    int32_t ky_end = args.KH;
+    const int32_t y_max = args.IH - 1 - base_y;
+    if (y_max < 0) {
+        ky_end = ky_start;
+    } else if (base_y + (args.KH - 1)*args.d1 >= args.IH) {
+        ky_end = min(ky_end, y_max/args.d1 + 1);
+    }
+
+    const uint64_t w_base   = (uint64_t) c * args.nb02;
+    const uint64_t src_base = (uint64_t) n * args.nb13 + (uint64_t) c * args.nb12;
+
+    const int32_t ow = (int32_t)(tgpig.x * ntg.x + tpitg.x);
+    if (ow >= args.OW) {
+        return;
+    }
+
+    float acc = 0.0f;
+
+    const int32_t base_x = ow*args.s0 - args.p0;
+
+    int32_t kx_start = 0;
+    if (base_x < 0) {
+        kx_start = (-base_x + args.d0 - 1)/args.d0;
+    }
+    int32_t kx_end = args.KW;
+    const int32_t x_max = args.IW - 1 - base_x;
+    if (x_max < 0) {
+        kx_end = kx_start;
+    } else if (base_x + (args.KW - 1)*args.d0 >= args.IW) {
+        kx_end = min(kx_end, x_max/args.d0 + 1);
+    }
+
+    if (ky_start < ky_end && kx_start < kx_end) {
+        for (int32_t ky = ky_start; ky < ky_end; ++ky) {
+            const int32_t iy = base_y + ky*args.d1;
+            const uint64_t src_row = src_base + (uint64_t) iy * args.nb11;
+            const uint64_t w_row = w_base + (uint64_t) ky * args.nb01;
+
+            for (int32_t kx = kx_start; kx < kx_end; ++kx) {
+                const int32_t ix = base_x + kx*args.d0;
+                const float x = *(device const float *)(src + src_row + (uint64_t) ix * args.nb10);
+                const float w = (float)(*(device const TK *)(weights + w_row + (uint64_t) kx * args.nb00));
+                acc += x * w;
+            }
+        }
+    }
+
+    const uint64_t dst_offs =
+        (uint64_t) n  * args.nb3 +
+        (uint64_t) c  * args.nb2 +
+        (uint64_t) oh * args.nb1 +
+        (uint64_t) ow * args.nb0;
+
+    *(device float *)(dst + dst_offs) = acc;
+}
+
+template [[host_name("kernel_conv_2d_dw_f32_f32")]]
+kernel void kernel_conv_2d_dw<float>(
+        constant ggml_metal_kargs_conv_2d_dw & args,
+        device const char * weights,
+        device const char * src,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        uint3   tpitg[[thread_position_in_threadgroup]],
+        uint3     ntg[[threads_per_threadgroup]]);
+
+template [[host_name("kernel_conv_2d_dw_f16_f32")]]
+kernel void kernel_conv_2d_dw<half>(
+        constant ggml_metal_kargs_conv_2d_dw & args,
+        device const char * weights,
+        device const char * src,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        uint3   tpitg[[thread_position_in_threadgroup]],
+        uint3     ntg[[threads_per_threadgroup]]);
+
+template [[host_name("kernel_conv_2d_dw_tiled_f32_f32")]]
+kernel void kernel_conv_2d_dw_tiled<float>(
+        constant ggml_metal_kargs_conv_2d_dw & args,
+        device const char * weights,
+        device const char * src,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        uint3   tpitg[[thread_position_in_threadgroup]],
+        uint3     ntg[[threads_per_threadgroup]]);
+
+template [[host_name("kernel_conv_2d_dw_tiled_f16_f32")]]
+kernel void kernel_conv_2d_dw_tiled<half>(
+        constant ggml_metal_kargs_conv_2d_dw & args,
+        device const char * weights,
+        device const char * src,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        uint3   tpitg[[thread_position_in_threadgroup]],
+        uint3     ntg[[threads_per_threadgroup]]);
+
+template <typename T>
+kernel void kernel_conv_3d(
+        constant ggml_metal_kargs_conv_3d & args,
+        device const  char * src0, // Weights [IC * OC, KD, KH, KW]
+        device const  char * src1, // Inputs  [IC * N,  ID, IH, IW]
+        device       char  * dst,  // Outputs [OC * N,  OD, OH, OW]
+        uint3 tgpig[[threadgroup_position_in_grid]],
+        uint3 tpitg[[thread_position_in_threadgroup]]) {
+
+    // 1. Un-flatten the spatial dimension from Grid X
+    int64_t spatial_idx = tgpig.x * 32 + tpitg.x;
+
+    if (spatial_idx >= args.OW * args.OH * args.OD) {
+        return; // Thread falls outside the spatial volume
+    }
+
+    int64_t od = spatial_idx / (args.OW * args.OH);
+    int64_t oh = (spatial_idx / args.OW) % args.OH;
+    int64_t ow = spatial_idx % args.OW;
+
+    // 2. Map Y to Channels, Z to Batch
+    int64_t oc = tgpig.y;
+    int64_t batch_idx = tgpig.z;
+
+    // 3. Calculate anchor coordinates in the Input volume
+    int64_t i_w_base = ow * args.s0 - args.p0;
+    int64_t i_h_base = oh * args.s1 - args.p1;
+    int64_t i_d_base = od * args.s2 - args.p2;
+
+    float sum = 0.0f;
+
+    // 4. Gather Loop (Iterate over Input Channels -> Depth -> Height -> Width)
+    for (int64_t ic = 0; ic < args.IC; ++ic) {
+
+        // ggml packs batch and channel together in the 4th dimension
+        int64_t src_cn_idx = batch_idx * args.IC + ic;
+        int64_t w_cn_idx   = oc * args.IC + ic;
+
+        for (int64_t kz = 0; kz < args.KD; ++kz) {
+            int64_t id = i_d_base + kz * args.d2;
+            if (id < 0 || id >= args.ID) continue; // Boundary check (Padding)
+
+            for (int64_t ky = 0; ky < args.KH; ++ky) {
+                int64_t ih = i_h_base + ky * args.d1;
+                if (ih < 0 || ih >= args.IH) continue;
+
+                for (int64_t kx = 0; kx < args.KW; ++kx) {
+                    int64_t iw = i_w_base + kx * args.d0;
+                    if (iw < 0 || iw >= args.IW) continue;
+
+                    // Convert multi-dimensional coordinates to flat byte offsets
+                    int64_t w_idx = kx*args.nb00 + ky*args.nb01 + kz*args.nb02 + w_cn_idx*args.nb03;
+                    int64_t i_idx = iw*args.nb10 + ih*args.nb11 + id*args.nb12 + src_cn_idx*args.nb13;
+
+                    // Dereference memory and cast weights to f32 if they were f16
+                    float w_val = (float)*(device const T*)((device const char*)src0 + w_idx);
+                    float i_val = *(device const float*)((device const char*)src1 + i_idx);
+
+                    sum += w_val * i_val;
+                }
+            }
+        }
+    }
+
+    // 5. Write the accumulated value out to RAM
+    int64_t dst_cn_idx = batch_idx * args.OC + oc;
+    int64_t d_idx = ow*args.nb0 + oh*args.nb1 + od*args.nb2 + dst_cn_idx*args.nb3;
+
+    *(device float*)(dst + d_idx) = sum;
+}
+
+// Explicit instantiations so the JIT compiler can find them by name
+template [[host_name("kernel_conv_3d_f32_f32")]]
+kernel void kernel_conv_3d<float>(
+    constant ggml_metal_kargs_conv_3d & args,
+    device const char * src0,
+    device const char * src1,
+    device       char  * dst,
+    uint3 tgpig[[threadgroup_position_in_grid]],
+    uint3 tpitg[[thread_position_in_threadgroup]]);
+
+// Explicit instantiation for f16 weights
+template [[host_name("kernel_conv_3d_f16_f32")]]
+kernel void kernel_conv_3d<half>(
+    constant ggml_metal_kargs_conv_3d & args,
+    device const char  * src0,
+    device const char * src1,
+    device       char  * dst,
+    uint3 tgpig[[threadgroup_position_in_grid]],
+    uint3 tpitg[[thread_position_in_threadgroup]]);

ggml/src/ggml-metal/kernels/dequantize.h

@@ -0,0 +1,686 @@
+#pragma once
+
+#include "common.h"
+
+#define GGML_COMMON_DECL_METAL
+#define GGML_COMMON_IMPL_METAL
+#if defined(GGML_METAL_EMBED_LIBRARY)
+__embed_ggml-common.h__
+#else
+#include "ggml-common.h"
+#endif
+
+#define QK_NL 16 // shared by mul_mm and get_rows_q instantiations
+
+// NOTE: this is not dequantizing - we are simply fitting the template
+template <typename type4x4>
+void dequantize_f32(device const float4x4 * src, short il, thread type4x4 & reg) {
+    reg = (type4x4)(*src);
+}
+
+template <typename type4>
+void dequantize_f32_t4(device const float4 * src, short il, thread type4 & reg) {
+    reg = (type4)(*src);
+}
+
+template <typename type4x4>
+void dequantize_f16(device const half4x4 * src, short il, thread type4x4 & reg) {
+    reg = (type4x4)(*src);
+}
+
+template <typename type4>
+void dequantize_f16_t4(device const half4 * src, short il, thread type4 & reg) {
+    reg = (type4)(*(src));
+}
+
+#if defined(GGML_METAL_HAS_BF16)
+template <typename type4x4>
+void dequantize_bf16(device const bfloat4x4 * src, short il, thread type4x4 & reg) {
+    reg = (type4x4)(*src);
+}
+
+template <typename type4>
+void dequantize_bf16_t4(device const bfloat4 * src, short il, thread type4 & reg) {
+    reg = (type4)(*(src));
+}
+#endif
+
+template <typename type4x4>
+void dequantize_q1_0(device const block_q1_0 * xb, short il, thread type4x4 & reg) {
+    device const uint8_t * qs = xb->qs;
+    const float d = xb->d;
+    const float neg_d = -d;
+
+    const int byte_offset = il * 2;  // il*16 bits = il*2 bytes
+    const uint8_t b0 = qs[byte_offset];
+    const uint8_t b1 = qs[byte_offset + 1];
+
+    float4x4 reg_f;
+
+    reg_f[0][0] = select(neg_d, d, bool(b0 & 0x01));
+    reg_f[0][1] = select(neg_d, d, bool(b0 & 0x02));
+    reg_f[0][2] = select(neg_d, d, bool(b0 & 0x04));
+    reg_f[0][3] = select(neg_d, d, bool(b0 & 0x08));
+    reg_f[1][0] = select(neg_d, d, bool(b0 & 0x10));
+    reg_f[1][1] = select(neg_d, d, bool(b0 & 0x20));
+    reg_f[1][2] = select(neg_d, d, bool(b0 & 0x40));
+    reg_f[1][3] = select(neg_d, d, bool(b0 & 0x80));
+
+    reg_f[2][0] = select(neg_d, d, bool(b1 & 0x01));
+    reg_f[2][1] = select(neg_d, d, bool(b1 & 0x02));
+    reg_f[2][2] = select(neg_d, d, bool(b1 & 0x04));
+    reg_f[2][3] = select(neg_d, d, bool(b1 & 0x08));
+    reg_f[3][0] = select(neg_d, d, bool(b1 & 0x10));
+    reg_f[3][1] = select(neg_d, d, bool(b1 & 0x20));
+    reg_f[3][2] = select(neg_d, d, bool(b1 & 0x40));
+    reg_f[3][3] = select(neg_d, d, bool(b1 & 0x80));
+
+    reg = (type4x4) reg_f;
+}
+
+template <typename type4>
+void dequantize_q1_0_t4(device const block_q1_0 * xb, short il, thread type4 & reg) {
+    const float d = xb->d;
+    const float neg_d = -d;
+    const int base = il * 4;
+    const uint8_t byte = xb->qs[base / 8];
+    const int s = base % 8;
+
+    float4 reg_f;
+    reg_f[0] = select(neg_d, d, bool((byte >> (s    )) & 1));
+    reg_f[1] = select(neg_d, d, bool((byte >> (s + 1)) & 1));
+    reg_f[2] = select(neg_d, d, bool((byte >> (s + 2)) & 1));
+    reg_f[3] = select(neg_d, d, bool((byte >> (s + 3)) & 1));
+
+    reg = (type4) reg_f;
+}
+
+template <typename type4x4>
+void dequantize_q4_0(device const block_q4_0 * xb, short il, thread type4x4 & reg) {
+    device const uint16_t * qs = ((device const uint16_t *)xb + 1);
+    const float d1 = il ? (xb->d / 16.h) : xb->d;
+    const float d2 = d1 / 256.f;
+    const float md = -8.h * xb->d;
+    const ushort mask0 = il ? 0x00F0 : 0x000F;
+    const ushort mask1 = mask0 << 8;
+
+    float4x4 reg_f;
+
+    for (int i = 0; i < 8; i++) {
+        reg_f[i/2][2*(i%2) + 0] = d1 * (qs[i] & mask0) + md;
+        reg_f[i/2][2*(i%2) + 1] = d2 * (qs[i] & mask1) + md;
+    }
+
+    reg = (type4x4) reg_f;
+}
+
+template <typename type4>
+void dequantize_q4_0_t4(device const block_q4_0 * xb, short il, thread type4 & reg) {
+    device const uint16_t * qs = ((device const uint16_t *)xb + 1);
+    const float d1 = (il/4) ? (xb->d / 16.h) : xb->d;
+    const float d2 = d1 / 256.f;
+    const float md = -8.h * xb->d;
+    const ushort mask0 = (il/4) ? 0x00F0 : 0x000F;
+    const ushort mask1 = mask0 << 8;
+
+    for (int i = 0; i < 2; i++) {
+        reg[2*i + 0] = d1 * (qs[2*(il%4) + i] & mask0) + md;
+        reg[2*i + 1] = d2 * (qs[2*(il%4) + i] & mask1) + md;
+    }
+}
+
+
+
+template <typename type4x4>
+void dequantize_q4_1(device const block_q4_1 * xb, short il, thread type4x4 & reg) {
+    device const uint16_t * qs = ((device const uint16_t *)xb + 2);
+    const float d1 = il ? (xb->d / 16.h) : xb->d;
+    const float d2 = d1 / 256.f;
+    const float  m = xb->m;
+    const ushort mask0 = il ? 0x00F0 : 0x000F;
+    const ushort mask1 = mask0 << 8;
+
+    float4x4 reg_f;
+
+    for (int i = 0; i < 8; i++) {
+        reg_f[i/2][2*(i%2) + 0] = ((qs[i] & mask0) * d1) + m;
+        reg_f[i/2][2*(i%2) + 1] = ((qs[i] & mask1) * d2) + m;
+    }
+
+    reg = (type4x4) reg_f;
+}
+
+template <typename type4>
+void dequantize_q4_1_t4(device const block_q4_1 * xb, short il, thread type4 & reg) {
+    device const uint16_t * qs = ((device const uint16_t *)xb + 2);
+    const float d1 = (il/4) ? (xb->d / 16.h) : xb->d;
+    const float d2 = d1 / 256.f;
+    const float  m = xb->m;
+    const ushort mask0 = (il/4) ? 0x00F0 : 0x000F;
+    const ushort mask1 = mask0 << 8;
+
+    for (int i = 0; i < 2; i++) {
+        reg[2*i + 0] = d1 * (qs[2*(il%4) + i] & mask0) + m;
+        reg[2*i + 1] = d2 * (qs[2*(il%4) + i] & mask1) + m;
+    }
+}
+
+template <typename type4x4>
+void dequantize_q5_0(device const block_q5_0 * xb, short il, thread type4x4 & reg) {
+    device const uint16_t * qs = ((device const uint16_t *)xb + 3);
+    const float d = xb->d;
+    const float md = -16.h * xb->d;
+    const ushort mask = il ? 0x00F0 : 0x000F;
+
+    const uint32_t qh = *((device const uint32_t *)xb->qh);
+
+    const int x_mv = il ? 4 : 0;
+
+    const int gh_mv = il ? 12 : 0;
+    const int gh_bk = il ?  0 : 4;
+
+    float4x4 reg_f;
+
+    for (int i = 0; i < 8; i++) {
+        // extract the 5-th bits for x0 and x1
+        const uint8_t xh_0 = ((qh >> (gh_mv + 2*i  )) << gh_bk) & 0x10;
+        const uint8_t xh_1 = ((qh >> (gh_mv + 2*i+1)) << gh_bk) & 0x10;
+
+        // combine the 4-bits from qs with the 5th bit
+        const int32_t x0 = ((((qs[i]     ) & mask) >> x_mv) | xh_0);
+        const int32_t x1 = ((((qs[i] >> 8) & mask) >> x_mv) | xh_1);
+
+        reg_f[i/2][2*(i%2) + 0] = d * x0 + md;
+        reg_f[i/2][2*(i%2) + 1] = d * x1 + md;
+    }
+
+    reg = (type4x4) reg_f;
+}
+
+template <typename type4>
+void dequantize_q5_0_t4(device const block_q5_0 * xb, short il, thread type4 & reg) {
+    device const uint16_t * qs = ((device const uint16_t *)xb + 3);
+    const float d = xb->d;
+    const float md = -16.h * xb->d;
+    const ushort mask = (il/4) ? 0x00F0 : 0x000F;
+
+    const uint32_t qh = *((device const uint32_t *)xb->qh);
+
+    const int x_mv = (il/4) ? 4 : 0;
+
+    const int gh_mv = (il/4) ? 12 : 0;
+    const int gh_bk = (il/4) ?  0 : 4;
+
+    for (int ii = 0; ii < 2; ii++) {
+        int i = 2*(il%4) + ii;
+
+        // extract the 5-th bits for x0 and x1
+        const uint8_t xh_0 = ((qh >> (gh_mv + 2*i  )) << gh_bk) & 0x10;
+        const uint8_t xh_1 = ((qh >> (gh_mv + 2*i+1)) << gh_bk) & 0x10;
+
+        // combine the 4-bits from qs with the 5th bit
+        const int32_t x0 = ((((qs[i]     ) & mask) >> x_mv) | xh_0);
+        const int32_t x1 = ((((qs[i] >> 8) & mask) >> x_mv) | xh_1);
+
+        reg[2*ii + 0] = d * x0 + md;
+        reg[2*ii + 1] = d * x1 + md;
+    }
+}
+
+template <typename type4x4>
+void dequantize_q5_1(device const block_q5_1 * xb, short il, thread type4x4 & reg) {
+    device const uint16_t * qs = ((device const uint16_t *)xb + 4);
+    const float d = xb->d;
+    const float m = xb->m;
+    const ushort mask = il ? 0x00F0 : 0x000F;
+
+    const uint32_t qh = *((device const uint32_t *)xb->qh);
+
+    const int x_mv = il ? 4 : 0;
+
+    const int gh_mv = il ? 12 : 0;
+    const int gh_bk = il ?  0 : 4;
+
+    float4x4 reg_f;
+
+    for (int i = 0; i < 8; i++) {
+        // extract the 5-th bits for x0 and x1
+        const uint8_t xh_0 = ((qh >> (gh_mv + 2*i  )) << gh_bk) & 0x10;
+        const uint8_t xh_1 = ((qh >> (gh_mv + 2*i+1)) << gh_bk) & 0x10;
+
+        // combine the 4-bits from qs with the 5th bit
+        const int32_t x0 = ((((qs[i]     ) & mask) >> x_mv) | xh_0);
+        const int32_t x1 = ((((qs[i] >> 8) & mask) >> x_mv) | xh_1);
+
+        reg_f[i/2][2*(i%2) + 0] = d * x0 + m;
+        reg_f[i/2][2*(i%2) + 1] = d * x1 + m;
+    }
+
+    reg = (type4x4) reg_f;
+}
+
+template <typename type4>
+void dequantize_q5_1_t4(device const block_q5_1 * xb, short il, thread type4 & reg) {
+    device const uint16_t * qs = ((device const uint16_t *)xb + 4);
+    const float d = xb->d;
+    const float m = xb->m;
+    const ushort mask = (il/4) ? 0x00F0 : 0x000F;
+
+    const uint32_t qh = *((device const uint32_t *)xb->qh);
+
+    const int x_mv = (il/4) ? 4 : 0;
+
+    const int gh_mv = (il/4) ? 12 : 0;
+    const int gh_bk = (il/4) ?  0 : 4;
+
+    for (int ii = 0; ii < 2; ii++) {
+        int i = 2*(il%4) + ii;
+
+        // extract the 5-th bits for x0 and x1
+        const uint8_t xh_0 = ((qh >> (gh_mv + 2*i  )) << gh_bk) & 0x10;
+        const uint8_t xh_1 = ((qh >> (gh_mv + 2*i+1)) << gh_bk) & 0x10;
+
+        // combine the 4-bits from qs with the 5th bit
+        const int32_t x0 = ((((qs[i]     ) & mask) >> x_mv) | xh_0);
+        const int32_t x1 = ((((qs[i] >> 8) & mask) >> x_mv) | xh_1);
+
+        reg[2*ii + 0] = d * x0 + m;
+        reg[2*ii + 1] = d * x1 + m;
+    }
+}
+
+template <typename type4x4>
+void dequantize_q8_0(device const block_q8_0 *xb, short il, thread type4x4 & reg) {
+    device const int8_t * qs = ((device const int8_t *)xb->qs);
+    const float d = xb->d;
+
+    float4x4 reg_f;
+
+    for (int i = 0; i < 16; i++) {
+        reg_f[i/4][i%4] = (qs[i + 16*il] * d);
+    }
+
+    reg = (type4x4) reg_f;
+}
+
+template <typename type4>
+void dequantize_q8_0_t4(device const block_q8_0 *xb, short il, thread type4 & reg) {
+    device const int8_t * qs = ((device const int8_t *)xb->qs);
+    const float d = xb->d;
+
+    for (int i = 0; i < 4; i++) {
+        reg[i] = (qs[4*(il%4) + i + 16*(il/4)] * d);
+    }
+}
+
+template <typename type4x4>
+void dequantize_mxfp4(device const block_mxfp4 * xb, short il, thread type4x4 & reg) {
+    device const uint8_t * q2 = (device const uint8_t *)xb->qs;
+
+    const float d = e8m0_to_fp32(xb->e);
+    const uint8_t shr = il >= 1 ? 4 : 0;
+
+    for (int i = 0; i < 4; ++i) {
+        reg[i][0] = d * kvalues_mxfp4_f[(q2[4*i + 0] >> shr) & 0x0F];
+        reg[i][1] = d * kvalues_mxfp4_f[(q2[4*i + 1] >> shr) & 0x0F];
+        reg[i][2] = d * kvalues_mxfp4_f[(q2[4*i + 2] >> shr) & 0x0F];
+        reg[i][3] = d * kvalues_mxfp4_f[(q2[4*i + 3] >> shr) & 0x0F];
+    }
+}
+
+template <typename type4>
+void dequantize_mxfp4_t4(device const block_mxfp4 * xb, short il, thread type4 & reg) {
+    device const uint8_t * q2 = (device const uint8_t *)xb->qs;
+
+    const float d = e8m0_to_fp32(xb->e);
+    const short il4 = il%4;
+
+    const uint8_t shr = il >= 4 ? 4 : 0;
+
+    reg[0] = d * kvalues_mxfp4_f[(q2[4*il4 + 0] >> shr) & 0x0F];
+    reg[1] = d * kvalues_mxfp4_f[(q2[4*il4 + 1] >> shr) & 0x0F];
+    reg[2] = d * kvalues_mxfp4_f[(q2[4*il4 + 2] >> shr) & 0x0F];
+    reg[3] = d * kvalues_mxfp4_f[(q2[4*il4 + 3] >> shr) & 0x0F];
+}
+
+template <typename type4x4>
+void dequantize_q2_K(device const block_q2_K *xb, short il, thread type4x4 & reg) {
+    const float d = xb->d;
+    const float min = xb->dmin;
+    device const uint8_t * q = (device const uint8_t *)xb->qs;
+    float dl, ml;
+    uint8_t sc = xb->scales[il];
+
+    q = q + 32*(il/8) + 16*(il&1);
+    il = (il/2)%4;
+
+    half  coef = il>1 ? (il>2 ? 1/64.h : 1/16.h) : (il>0 ? 1/4.h : 1.h);
+    uchar mask = il>1 ? (il>2 ? 192    : 48)     : (il>0 ? 12    : 3);
+    dl = d * (sc & 0xF) * coef, ml = min * (sc >> 4);
+    for (int i = 0; i < 16; ++i) {
+        reg[i/4][i%4] = dl * (q[i] & mask) - ml;
+    }
+}
+
+template <typename type4x4>
+void dequantize_q3_K(device const block_q3_K *xb, short il, thread type4x4 & reg) {
+    const half d_all = xb->d;
+    device const uint8_t * q = (device const uint8_t *)xb->qs;
+    device const uint8_t * h = (device const uint8_t *)xb->hmask;
+    device const int8_t * scales = (device const int8_t *)xb->scales;
+
+    q = q + 32 * (il/8) + 16 * (il&1);
+    h = h + 16 * (il&1);
+    uint8_t m = 1 << (il/2);
+    uint16_t kmask1 = (il/4)>1 ? ((il/4)>2 ? 192 : 48) : \
+                                 ((il/4)>0 ? 12  : 3);
+    uint16_t kmask2 = il/8 ? 0xF0 : 0x0F;
+    uint16_t scale_2 = scales[il%8], scale_1 = scales[8 + il%4];
+    int16_t  dl_int = (il/4)&1 ? (scale_2&kmask2) | ((scale_1&kmask1) << 2)
+                               : (scale_2&kmask2) | ((scale_1&kmask1) << 4);
+    float dl = il<8 ? d_all * (dl_int - 32.f) : d_all * (dl_int / 16.f - 32.f);
+    const float ml = 4.f * dl;
+
+    il = (il/2) & 3;
+    const half    coef = il>1 ? (il>2 ? 1/64.h : 1/16.h) : (il>0 ? 1/4.h : 1.h);
+    const uint8_t mask = il>1 ? (il>2 ? 192    : 48)     : (il>0 ? 12    : 3);
+    dl *= coef;
+
+    for (int i = 0; i < 16; ++i) {
+        reg[i/4][i%4] = dl * (q[i] & mask) - (h[i] & m ? 0 : ml);
+    }
+}
+
+static inline uchar2 get_scale_min_k4_just2(int j, int k, device const uchar * q) {
+    return j < 4 ? uchar2{uchar(q[j+0+k] & 63), uchar(q[j+4+k] & 63)}
+                 : uchar2{uchar((q[j+4+k] & 0xF) | ((q[j-4+k] & 0xc0) >> 2)), uchar((q[j+4+k] >> 4) | ((q[j-0+k] & 0xc0) >> 2))};
+}
+
+template <typename type4x4>
+void dequantize_q4_K(device const block_q4_K * xb, short il, thread type4x4 & reg) {
+    device const uchar * q = xb->qs;
+
+    short is = (il/4) * 2;
+    q = q + (il/4) * 32 + 16 * (il&1);
+    il = il & 3;
+    const uchar2 sc = get_scale_min_k4_just2(is, il/2, xb->scales);
+    const float d   = il < 2 ? xb->d : xb->d / 16.h;
+    const float min = xb->dmin;
+    const float dl = d * sc[0];
+    const float ml = min * sc[1];
+
+    const ushort mask = il < 2 ? 0x0F : 0xF0;
+    for (int i = 0; i < 16; ++i) {
+        reg[i/4][i%4] = dl * (q[i] & mask) - ml;
+    }
+}
+
+template <typename type4x4>
+void dequantize_q5_K(device const block_q5_K *xb, short il, thread type4x4 & reg) {
+    device const uint8_t * q  = xb->qs;
+    device const uint8_t * qh = xb->qh;
+
+    short is = (il/4) * 2;
+    q  = q + 32 * (il/4) + 16 * (il&1);
+    qh = qh + 16 * (il&1);
+    uint8_t ul = 1 << (il/2);
+    il = il & 3;
+    const uchar2 sc = get_scale_min_k4_just2(is, il/2, xb->scales);
+    const float d = il < 2 ? xb->d : xb->d / 16.f;
+    const float min = xb->dmin;
+    const float dl = d * sc[0];
+    const float ml = min * sc[1];
+
+    const ushort mask  = il<2 ? 0x0F : 0xF0;
+    const float qh_val = il<2 ? 16.f : 256.f;
+    for (int i = 0; i < 16; ++i) {
+        reg[i/4][i%4] = dl * ((q[i] & mask) + (qh[i] & ul ? qh_val : 0)) - ml;
+    }
+}
+
+template <typename type4x4>
+void dequantize_q6_K(device const block_q6_K *xb, short il, thread type4x4 & reg) {
+    const half d_all = xb->d;
+    device const uint16_t * ql = (device const uint16_t *)xb->ql;
+    device const uint16_t * qh = (device const uint16_t *)xb->qh;
+    device const int8_t * scales = (device const int8_t *)xb->scales;
+
+    ql = ql + 32*(il/8) + 16*((il/2)&1) + 8*(il&1);
+    qh = qh + 16*(il/8) + 8*(il&1);
+    float sc = scales[(il%2) + 2 * ((il/2))];
+    il = (il/2) & 3;
+
+    const uint32_t kmask1 = il>1 ? (il>2 ? 0xC0C0C0C0 : 0x30303030) : (il>0 ? 0x0C0C0C0C : 0x03030303);
+    const uint32_t kmask2 = il>1 ? 0xF0F0F0F0                       : 0x0F0F0F0F;
+    const float ml = d_all * sc * 32.f;
+    const float dl0 = d_all * sc;
+    const float dl1 = dl0 / 256.f;
+    const float dl2 = dl0 / (256.f * 256.f);
+    const float dl3 = dl0 / (256.f * 256.f * 256.f);
+    const uint8_t shr_h = il>2 ? 2 : 0;
+    const uint8_t shl_h = il>1 ? 0 : (il>0 ? 2 : 4);
+    const uint8_t shr_l = il>1 ? 4 : 0;
+    for (int i = 0; i < 4; ++i) {
+        const uint32_t  low = (ql[2*i] | (uint32_t)(ql[2*i+1] << 16)) & kmask2;
+        const uint32_t high = (qh[2*i] | (uint32_t)(qh[2*i+1] << 16)) & kmask1;
+        const uint32_t q = ((high << shl_h) >> shr_h) | (low >> shr_l);
+        reg[i][0] = dl0 *  ((half)(q & 0xFF))       - ml;
+        reg[i][1] = dl1 * ((float)(q & 0xFF00))     - ml;
+        reg[i][2] = dl2 * ((float)(q & 0xFF0000))   - ml;
+        reg[i][3] = dl3 * ((float)(q & 0xFF000000)) - ml;
+    }
+}
+
+template <typename type4x4>
+void dequantize_iq2_xxs(device const block_iq2_xxs * xb, short il, thread type4x4 & reg) {
+    // il is 0...15 for QK_K = 256 => index of block of 32 is il/2
+    const float d = xb->d;
+    const int ib32 = il/2;
+    il = il%2;
+    // il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
+    // each block of 32 needs 2 uint32_t's for the quants & scale, so 4 uint16_t's.
+    device const uint16_t * q2 = xb->qs + 4*ib32;
+    const uint32_t aux32_g = q2[0] | (q2[1] << 16);
+    const uint32_t aux32_s = q2[2] | (q2[3] << 16);
+    thread const uint8_t * aux8 = (thread const uint8_t *)&aux32_g;
+    const float dl = d * (0.5f + (aux32_s >> 28)) * 0.25f;
+    constant uint8_t * grid = (constant uint8_t *)(iq2xxs_grid + aux8[2*il+0]);
+    uint8_t signs = ksigns_iq2xs[(aux32_s >> 14*il) & 127];
+    for (int i = 0; i < 8; ++i) {
+        reg[i/4][i%4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f);
+    }
+    grid = (constant uint8_t *)(iq2xxs_grid + aux8[2*il+1]);
+    signs = ksigns_iq2xs[(aux32_s >> (14*il+7)) & 127];
+    for (int i = 0; i < 8; ++i) {
+        reg[2+i/4][i%4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f);
+    }
+}
+
+template <typename type4x4>
+void dequantize_iq2_xs(device const block_iq2_xs * xb, short il, thread type4x4 & reg) {
+    // il is 0...15 for QK_K = 256 => index of block of 32 is il/2
+    const float d = xb->d;
+    const int ib32 = il/2;
+    il = il%2;
+    // il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
+    device const uint16_t * q2 = xb->qs + 4*ib32;
+    const float dl = d * (0.5f + ((xb->scales[ib32] >> 4*il) & 0xf)) * 0.25f;
+    constant uint8_t * grid = (constant uint8_t *)(iq2xs_grid + (q2[2*il+0] & 511));
+    uint8_t signs = ksigns_iq2xs[q2[2*il+0] >> 9];
+    for (int i = 0; i < 8; ++i) {
+        reg[i/4][i%4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f);
+    }
+    grid = (constant uint8_t *)(iq2xs_grid + (q2[2*il+1] & 511));
+    signs = ksigns_iq2xs[q2[2*il+1] >> 9];
+    for (int i = 0; i < 8; ++i) {
+        reg[2+i/4][i%4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f);
+    }
+}
+
+template <typename type4x4>
+void dequantize_iq3_xxs(device const block_iq3_xxs * xb, short il, thread type4x4 & reg) {
+    // il is 0...15 for QK_K = 256 => index of block of 32 is il/2
+    const float d = xb->d;
+    const int ib32 = il/2;
+    il = il%2;
+    // il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
+    device const uint8_t * q3 = xb->qs + 8*ib32;
+    device const uint16_t * gas = (device const uint16_t *)(xb->qs + QK_K/4) + 2*ib32;
+    const uint32_t aux32 = gas[0] | (gas[1] << 16);
+    const float dl = d * (0.5f + (aux32 >> 28)) * 0.5f;
+    constant uint8_t * grid1 = (constant uint8_t *)(iq3xxs_grid + q3[4*il+0]);
+    constant uint8_t * grid2 = (constant uint8_t *)(iq3xxs_grid + q3[4*il+1]);
+    uint8_t signs = ksigns_iq2xs[(aux32 >> 14*il) & 127];
+    for (int i = 0; i < 4; ++i) {
+        reg[0][i] = dl * grid1[i] * (signs & kmask_iq2xs[i+0] ? -1.f : 1.f);
+        reg[1][i] = dl * grid2[i] * (signs & kmask_iq2xs[i+4] ? -1.f : 1.f);
+    }
+    grid1 = (constant uint8_t *)(iq3xxs_grid + q3[4*il+2]);
+    grid2 = (constant uint8_t *)(iq3xxs_grid + q3[4*il+3]);
+    signs = ksigns_iq2xs[(aux32 >> (14*il+7)) & 127];
+    for (int i = 0; i < 4; ++i) {
+        reg[2][i] = dl * grid1[i] * (signs & kmask_iq2xs[i+0] ? -1.f : 1.f);
+        reg[3][i] = dl * grid2[i] * (signs & kmask_iq2xs[i+4] ? -1.f : 1.f);
+    }
+}
+
+template <typename type4x4>
+void dequantize_iq3_s(device const block_iq3_s * xb, short il, thread type4x4 & reg) {
+    // il is 0...15 for QK_K = 256 => index of block of 32 is il/2
+    const float d = xb->d;
+    const int ib32 = il/2;
+    il = il%2;
+    // il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
+    device const uint8_t * qs = xb->qs + 8*ib32;
+    device const uint8_t * signs = xb->signs + 4*ib32 + 2*il;
+    const uint8_t qh = xb->qh[ib32] >> 4*il;
+    const float dl = d * (1 + 2*((xb->scales[ib32/2] >> 4*(ib32%2)) & 0xf));
+    constant uint8_t * grid1 = (constant uint8_t *)(iq3s_grid + (qs[4*il+0] | ((qh << 8) & 256)));
+    constant uint8_t * grid2 = (constant uint8_t *)(iq3s_grid + (qs[4*il+1] | ((qh << 7) & 256)));
+    for (int i = 0; i < 4; ++i) {
+        reg[0][i] = dl * grid1[i] * select(1, -1, signs[0] & kmask_iq2xs[i+0]);
+        reg[1][i] = dl * grid2[i] * select(1, -1, signs[0] & kmask_iq2xs[i+4]);
+    }
+    grid1 = (constant uint8_t *)(iq3s_grid + (qs[4*il+2] | ((qh << 6) & 256)));
+    grid2 = (constant uint8_t *)(iq3s_grid + (qs[4*il+3] | ((qh << 5) & 256)));
+    for (int i = 0; i < 4; ++i) {
+        reg[2][i] = dl * grid1[i] * select(1, -1, signs[1] & kmask_iq2xs[i+0]);
+        reg[3][i] = dl * grid2[i] * select(1, -1, signs[1] & kmask_iq2xs[i+4]);
+    }
+}
+
+template <typename type4x4>
+void dequantize_iq2_s(device const block_iq2_s * xb, short il, thread type4x4 & reg) {
+    // il is 0...15 for QK_K = 256 => index of block of 32 is il/2
+    const float d = xb->d;
+    const int ib32 = il/2;
+    il = il%2;
+    // il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
+    device const uint8_t * qs = xb->qs + 4*ib32 + 2*il;
+    device const uint8_t * signs = qs + QK_K/8;
+    const uint8_t qh = xb->qh[ib32] >> 4*il;
+    const float dl = d * (0.5f + ((xb->scales[ib32] >> 4*il) & 0xf)) * 0.25f;
+    constant uint8_t * grid1 = (constant uint8_t *)(iq2s_grid + (qs[0] | ((qh << 8) & 0x300)));
+    constant uint8_t * grid2 = (constant uint8_t *)(iq2s_grid + (qs[1] | ((qh << 6) & 0x300)));
+    for (int i = 0; i < 8; ++i) {
+        reg[i/4+0][i%4] = dl * grid1[i] * select(1, -1, signs[0] & kmask_iq2xs[i]);
+        reg[i/4+2][i%4] = dl * grid2[i] * select(1, -1, signs[1] & kmask_iq2xs[i]);
+    }
+}
+
+template <typename type4x4>
+void dequantize_iq1_s(device const block_iq1_s * xb, short il, thread type4x4 & reg) {
+    // il is 0...15 for QK_K = 256 => index of block of 32 is il/2
+    const int ib32 = il/2;
+    il = il%2;
+    const float d = xb->d;
+    device const uint8_t  * qs = xb->qs + 4*ib32 + 2*il;
+    device const uint16_t * qh = xb->qh;
+    const float dl = d * (2*((qh[ib32] >> 12) & 7) + 1);
+    const float ml = dl * (qh[ib32] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA);
+    const uint16_t h = qh[ib32] >> 6*il;
+    constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((h << 8) & 0x700)));
+    constant uint8_t * grid2 = (constant uint8_t *)(iq1s_grid_gpu + (qs[1] | ((h << 5) & 0x700)));
+    for (int i = 0; i < 4; ++i) {
+        reg[0][i] = dl * (grid1[i] & 0xf) + ml;
+        reg[1][i] = dl * (grid1[i] >>  4) + ml;
+        reg[2][i] = dl * (grid2[i] & 0xf) + ml;
+        reg[3][i] = dl * (grid2[i] >>  4) + ml;
+    }
+}
+
+template <typename type4x4>
+void dequantize_iq1_m(device const block_iq1_m * xb, short il, thread type4x4 & reg) {
+    // il is 0...15 for QK_K = 256 => index of block of 32 is il/2
+    const int ib32 = il/2;
+    il = il%2;
+    device const uint16_t * sc = (device const uint16_t *)xb->scales;
+
+    iq1m_scale_t scale;
+    scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
+    const float d = scale.f16;
+
+    device const uint8_t * qs = xb->qs + 4*ib32 + 2*il;
+    device const uint8_t * qh = xb->qh + 2*ib32 + il;
+
+    const float dl  = d * (2*((sc[ib32/2] >> (6*(ib32%2)+3*il)) & 7) + 1);
+    const float ml1 = dl * (qh[0] & 0x08 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA);
+    const float ml2 = dl * (qh[0] & 0x80 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA);
+    constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((qh[0] << 8) & 0x700)));
+    constant uint8_t * grid2 = (constant uint8_t *)(iq1s_grid_gpu + (qs[1] | ((qh[0] << 4) & 0x700)));
+    for (int i = 0; i < 4; ++i) {
+        reg[0][i] = dl * (grid1[i] & 0xf) + ml1;
+        reg[1][i] = dl * (grid1[i] >>  4) + ml1;
+        reg[2][i] = dl * (grid2[i] & 0xf) + ml2;
+        reg[3][i] = dl * (grid2[i] >>  4) + ml2;
+    }
+}
+
+template <typename type4x4>
+void dequantize_iq4_nl(device const block_iq4_nl * xb, short il, thread type4x4 & reg) {
+    device const uint16_t * q4 = (device const uint16_t *)xb->qs;
+    const float d = xb->d;
+    uint32_t aux32;
+    thread const uint8_t * q8 = (thread const uint8_t *)&aux32;
+    for (int i = 0; i < 4; ++i) {
+        aux32 = ((q4[2*i] | (q4[2*i+1] << 16)) >> 4*il) & 0x0f0f0f0f;
+        reg[i][0] = d * kvalues_iq4nl_f[q8[0]];
+        reg[i][1] = d * kvalues_iq4nl_f[q8[1]];
+        reg[i][2] = d * kvalues_iq4nl_f[q8[2]];
+        reg[i][3] = d * kvalues_iq4nl_f[q8[3]];
+    }
+}
+
+template <typename type4>
+void dequantize_iq4_nl_t4(device const block_iq4_nl * xb, short il, thread type4 & reg) {
+    device const uint16_t * q4 = (device const uint16_t *)xb->qs;
+    const float d = xb->d;
+    uint32_t aux32;
+    thread const uint8_t * q8 = (thread const uint8_t *)&aux32;
+    aux32 = ((q4[2*(il%4)] | (q4[2*(il%4)+1] << 16)) >> 4*(il/4)) & 0x0f0f0f0f;
+    reg[0] = d * kvalues_iq4nl_f[q8[0]];
+    reg[1] = d * kvalues_iq4nl_f[q8[1]];
+    reg[2] = d * kvalues_iq4nl_f[q8[2]];
+    reg[3] = d * kvalues_iq4nl_f[q8[3]];
+}
+
+template <typename type4x4>
+void dequantize_iq4_xs(device const block_iq4_xs * xb, short il, thread type4x4 & reg) {
+    // il is 0...15 for QK_K = 256 => index of block of 32 is il/2
+    const int ib32 = il/2;
+    il = il%2;
+    // il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
+    device const uint32_t * q4 = (device const uint32_t *)xb->qs + 4*ib32;
+    const int ls = ((xb->scales_l[ib32/2] >> 4*(ib32%2)) & 0xf) | (((xb->scales_h >> 2*ib32) & 3) << 4);
+    const float d = (float)xb->d * (ls - 32);
+    uint32_t aux32;
+    thread const uint8_t * q8 = (thread const uint8_t *)&aux32;
+    for (int i = 0; i < 4; ++i) {
+        aux32 = (q4[i] >> 4*il) & 0x0f0f0f0f;
+        reg[i][0] = d * kvalues_iq4nl_f[q8[0]];
+        reg[i][1] = d * kvalues_iq4nl_f[q8[1]];
+        reg[i][2] = d * kvalues_iq4nl_f[q8[2]];
+        reg[i][3] = d * kvalues_iq4nl_f[q8[3]];
+    }
+}
+

ggml/src/ggml-metal/kernels/gated_delta_net.metal

@@ -0,0 +1,250 @@
+#include "common.h"
+
+constant short FC_gated_delta_net_ne20 [[function_constant(FC_GATED_DELTA_NET + 0)]];
+constant short FC_gated_delta_net_ne30 [[function_constant(FC_GATED_DELTA_NET + 1)]];
+constant short FC_gated_delta_net_K    [[function_constant(FC_GATED_DELTA_NET + 2)]];
+
+#if 1
+template<short NSG>
+kernel void kernel_gated_delta_net_impl(
+        constant ggml_metal_kargs_gated_delta_net & args,
+        device const char * q,
+        device const char * k,
+        device const char * v,
+        device const char * g,
+        device const char * b,
+        device const char * s,
+        device       char * dst,
+        uint3 tgpig[[threadgroup_position_in_grid]],
+        uint3 tpitg[[thread_position_in_threadgroup]],
+        uint3   ntg[[threads_per_threadgroup]])  {
+#define S_v FC_gated_delta_net_ne20
+#define G   FC_gated_delta_net_ne30
+#define K   FC_gated_delta_net_K
+
+    const uint tx = tpitg.x;
+    const uint ty = tpitg.y;
+
+    const uint i23 = tgpig.z; // B (n_seqs)
+    const uint i21 = tgpig.y; // H (head)
+    const uint i20 = tgpig.x*NSG + ty; // row within S_v
+
+    const uint i01 = i21 % args.ne01;
+    const uint i11 = i21 % args.ne11;
+
+    const float scale = 1.0f / sqrt((float)S_v);
+
+    // input state layout [S_v, S_v, H, n_seqs] (s0 only): per-seq stride is H*D.
+    // state is stored transposed: M[i20][is] = S[is][i20], so row i20 is contiguous
+    const uint state_in_base = (i23*args.ne21 + i21)*S_v*S_v + i20*S_v;
+    device const float * s_ptr = (device const float *) (s) + state_in_base;
+
+    float ls[NSG];
+
+    FOR_UNROLL (short j = 0; j < NSG; j++) {
+        const short is = tx*NSG + j;
+        ls[j] = s_ptr[is];
+    }
+
+    device float * dst_attn = (device float *) (dst) + (i23*args.ne22*args.ne21 + i21)*S_v + i20;
+
+    device const float * q_ptr = (device const float *) (q + i23*args.nb03 + i01*args.nb01);
+    device const float * k_ptr = (device const float *) (k + i23*args.nb13 + i11*args.nb11);
+    device const float * v_ptr = (device const float *) (v + i23*args.nb23 + i21*args.nb21);
+
+    device const float * b_ptr = (device const float *) (b) + (i23*args.ne22*args.ne21 + i21);
+    device const float * g_ptr = (device const float *) (g) + (i23*args.ne22*args.ne21 + i21)*G;
+
+    // snapshot slot mapping: slot 0 = most recent state, slot s = s tokens back.
+    // When n_tokens < K, only slots 0..n_tokens-1 are written; older slots are caller-owned.
+
+    // output state base offset: after attention scores
+    const uint attn_size = args.ne22 * args.ne21 * S_v * args.ne23;
+    // output state per-slot size: S_v * S_v * H * n_seqs
+    const uint state_size_per_snap = S_v * S_v * args.ne21 * args.ne23;
+    // per-(seq,head) offset within a slot
+    const uint state_out_base = (i23*args.ne21 + i21)*S_v*S_v + i20*S_v;
+
+    for (short t = 0; t < args.ne22; t++) {
+        float s_k = 0.0f;
+
+        if (G == 1) {
+            const float g_exp = exp(g_ptr[0]);
+
+            FOR_UNROLL (short j = 0; j < NSG; j++) {
+                const short is = tx*NSG + j;
+                ls[j] *= g_exp;
+
+                s_k += ls[j]*k_ptr[is];
+            }
+        } else {
+            // KDA
+            FOR_UNROLL (short j = 0; j < NSG; j++) {
+                const short is = tx*NSG + j;
+                ls[j] *= exp(g_ptr[is]);
+
+                s_k += ls[j]*k_ptr[is];
+            }
+        }
+
+        s_k = simd_sum(s_k);
+
+        const float d = (v_ptr[i20] - s_k)*b_ptr[0];
+
+        float y = 0.0f;
+
+        FOR_UNROLL (short j = 0; j < NSG; j++) {
+            const short is = tx*NSG + j;
+            ls[j] += k_ptr[is]*d;
+
+            y += ls[j]*q_ptr[is];
+        }
+
+        y = simd_sum(y);
+
+        if (tx == 0) {
+            dst_attn[t*args.ne21*S_v] = y*scale;
+        }
+
+        q_ptr += args.ns02;
+        k_ptr += args.ns12;
+        v_ptr += args.ns22;
+
+        b_ptr += args.ne21;
+        g_ptr += args.ne21*G;
+
+        if (K > 1) {
+            const int target_slot = (int)args.ne22 - 1 - (int)t;
+            if (target_slot >= 0 && target_slot < (int)K) {
+                device float * dst_state = (device float *) (dst) + attn_size + (uint)target_slot * state_size_per_snap + state_out_base;
+                FOR_UNROLL (short j = 0; j < NSG; j++) {
+                    const short is = tx*NSG + j;
+                    dst_state[is] = ls[j];
+                }
+            }
+        }
+    }
+
+    if (K == 1) {
+        device float * dst_state = (device float *) (dst) + attn_size + state_out_base;
+        FOR_UNROLL (short j = 0; j < NSG; j++) {
+            const short is = tx*NSG + j;
+            dst_state[is] = ls[j];
+        }
+    }
+
+#undef S_v
+#undef G
+#undef K
+}
+
+typedef decltype(kernel_gated_delta_net_impl<4>) kernel_gated_delta_net_t;
+
+template [[host_name("kernel_gated_delta_net_f32_1")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<1>;
+template [[host_name("kernel_gated_delta_net_f32_2")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<2>;
+template [[host_name("kernel_gated_delta_net_f32_4")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<4>;
+
+#else
+// a simplified version of the above
+// no performance improvement, so keep the above version for now
+
+template<typename T, short NSG>
+kernel void kernel_gated_delta_net_impl(
+        constant ggml_metal_kargs_gated_delta_net & args,
+        device const char * q,
+        device const char * k,
+        device const char * v,
+        device const char * g,
+        device const char * b,
+        device const char * s,
+        device       char * dst,
+        uint3 tgpig[[threadgroup_position_in_grid]],
+        uint3 tpitg[[thread_position_in_threadgroup]],
+        uint3   ntg[[threads_per_threadgroup]])  {
+#define S_v FC_gated_delta_net_ne20
+#define G   FC_gated_delta_net_ne30
+
+    const uint tx = tpitg.x;
+    const uint ty = tpitg.y;
+
+    const uint i23 = tgpig.z; // B
+    const uint i21 = tgpig.y; // H
+    const uint i20 = tgpig.x*NSG + ty;
+
+    const uint i01 = i21 % args.ne01;
+    const uint i11 = i21 % args.ne11;
+
+    const float scale = 1.0f / sqrt((float)S_v);
+
+    device const float * s_ptr = (device const float *) (s) + (i23*args.ne21 + i21)*S_v*S_v + i20;
+
+    float lsf[NSG];
+
+    FOR_UNROLL (short j = 0; j < NSG; j++) {
+        const short is = tx*NSG + j;
+        lsf[j] = s_ptr[is*S_v];
+    }
+
+    thread T * ls = (thread T *) (lsf);
+
+    device float * dst_attn = (device float *) (dst) + (i23*args.ne22*args.ne21 + i21)*S_v + i20;
+
+    device const float * q_ptr = (device const float *) (q + i23*args.nb03 + i01*args.nb01);
+    device const float * k_ptr = (device const float *) (k + i23*args.nb13 + i11*args.nb11);
+    device const float * v_ptr = (device const float *) (v + i23*args.nb23 + i21*args.nb21);
+
+    device const float * b_ptr  = (device const float *) (b) + (i23*args.ne22*args.ne21 + i21);
+    device const float * g_ptr  = (device const float *) (g) + (i23*args.ne22*args.ne21 + i21)*G;
+
+    for (short t = 0; t < args.ne22; t++) {
+        device const T * qt_ptr = (device const T *) (q_ptr);
+        device const T * kt_ptr = (device const T *) (k_ptr);
+        device const T * gt_ptr = (device const T *) (g_ptr);
+
+        if (G == 1) {
+            *ls *= exp(g_ptr[0]);
+        } else {
+            // KDA
+            *ls *= exp(gt_ptr[tx]);
+        }
+
+        const float s_k = simd_sum(dot(*ls, kt_ptr[tx]));
+
+        const float d = (v_ptr[i20] - s_k)*b_ptr[0];
+
+        *ls += kt_ptr[tx]*d;
+
+        const float y = simd_sum(dot(*ls, qt_ptr[tx]));
+
+        if (tx == 0) {
+            *dst_attn = y*scale;
+        }
+
+        q_ptr += args.ns02;
+        k_ptr += args.ns12;
+        v_ptr += args.ns22;
+
+        b_ptr += args.ne21;
+        g_ptr += args.ne21*G;
+
+        dst_attn += args.ne21*S_v;
+    }
+
+    device float * dst_state  = (device float *) (dst) + args.ne23*args.ne22*args.ne21*S_v + (i23*args.ne21 + i21)*S_v*S_v + i20;
+    device T     * dstt_state = (device T     *) (dst_state);
+
+    FOR_UNROLL (short j = 0; j < NSG; j++) {
+        const short is = tx*NSG + j;
+        dst_state[is*S_v] = lsf[j];
+    }
+
+#undef S_v
+#undef G
+}
+
+typedef decltype(kernel_gated_delta_net_impl<float4, 4>) kernel_gated_delta_net_t;
+
+template [[host_name("kernel_gated_delta_net_f32_1")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<float,  1>;
+template [[host_name("kernel_gated_delta_net_f32_2")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<float2, 2>;
+template [[host_name("kernel_gated_delta_net_f32_4")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<float4, 4>;
+#endif

ggml/src/ggml-metal/kernels/misc.metal

@@ -0,0 +1,347 @@
+#include "common.h"
+
+kernel void kernel_argmax_f32(
+        constant ggml_metal_kargs_argmax & args,
+        device   const char * src0,
+        device         char * dst,
+        threadgroup    char * shmem [[threadgroup(0)]],
+        uint  tgpig[[threadgroup_position_in_grid]],
+        uint  tpitg[[thread_position_in_threadgroup]],
+        uint  sgitg[[simdgroup_index_in_threadgroup]],
+        uint  tiisg[[thread_index_in_simdgroup]],
+        uint    ntg[[threads_per_threadgroup]]) {
+    device const float * x_row = (device const float *) ((device const char *) src0 + tgpig * args.nb01);
+
+    float   lmax = -INFINITY;
+    int32_t larg = -1;
+
+    for (int i00 = tpitg; i00 < args.ne00; i00 += ntg) {
+        if (x_row[i00] > lmax) {
+            lmax = x_row[i00];
+            larg = i00;
+        }
+    }
+
+    // find the argmax value in the block
+    float max_val = simd_max(lmax);
+    int32_t arg_val = simd_max(select(-1, larg, lmax == max_val));
+
+    device int32_t * dst_i32 = (device int32_t *) dst;
+
+    threadgroup   float * shared_maxval = (threadgroup   float *) shmem;
+    threadgroup int32_t * shared_argmax = (threadgroup int32_t *) shmem + N_SIMDWIDTH;
+
+    if (ntg > N_SIMDWIDTH) {
+        if (sgitg == 0) {
+            shared_maxval[tiisg] = -INFINITY;
+            shared_argmax[tiisg] = -1;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiisg == 0) {
+            shared_maxval[sgitg] = max_val;
+            shared_argmax[sgitg] = arg_val;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        max_val = shared_maxval[tiisg];
+        arg_val = shared_argmax[tiisg];
+
+        float max_val_reduced   = simd_max(max_val);
+        int32_t arg_val_reduced = simd_max(select(-1, arg_val, max_val == max_val_reduced));
+
+        dst_i32[tgpig] = arg_val_reduced;
+
+        return;
+    }
+
+    dst_i32[tgpig] = arg_val;
+}
+
+kernel void kernel_diag_f32(
+        constant ggml_metal_kargs_diag & args,
+        device   const char * src0,
+        device         char * dst,
+        uint3  tgpig[[threadgroup_position_in_grid]],
+        ushort tiitg[[thread_index_in_threadgroup]]) {
+    constexpr short NW = N_SIMDWIDTH;
+
+    const int32_t i3 = tgpig.z;
+    const int32_t i2 = tgpig.y;
+    const int32_t i1 = tgpig.x;
+
+    device const float * src0_ptr = (device const float *)(src0 +                i2*args.nb02 + i3*args.nb03);
+    device       float * dst_ptr  = (device       float *)(dst  + i1*args.nb01 + i2*args.nb2  + i3*args.nb3);
+
+    for (int i0 = tiitg; i0 < args.ne0; i0 += NW) {
+        dst_ptr[i0] = i0 == i1 ? src0_ptr[i0] : 0.0f;
+    }
+}
+
+kernel void kernel_roll_f32(
+    constant ggml_metal_kargs_roll & args,
+    device  const char * src0,
+    device        char * dst,
+    uint3 tgpig[[threadgroup_position_in_grid]],
+    uint3 tpitg[[thread_position_in_threadgroup]],
+    uint3   ntg[[threads_per_threadgroup]]) {
+
+    const int64_t i3 = tgpig.z;
+    const int64_t i2 = tgpig.y;
+    const int64_t i1 = tgpig.x;
+
+    device const float * src0_ptr = (device const float *) src0;
+    device       float * dst_ptr  = (device       float *) dst;
+
+    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
+        // apply shifts and wrap around
+        int64_t i00 = i0 - args.s0;
+        int64_t i01 = i1 - args.s1;
+        int64_t i02 = i2 - args.s2;
+        int64_t i03 = i3 - args.s3;
+
+        if (i00 < 0) { i00 += args.ne00; } else if (i00 >= args.ne00) { i00 -= args.ne00; }
+        if (i01 < 0) { i01 += args.ne01; } else if (i01 >= args.ne01) { i01 -= args.ne01; }
+        if (i02 < 0) { i02 += args.ne02; } else if (i02 >= args.ne02) { i02 -= args.ne02; }
+        if (i03 < 0) { i03 += args.ne03; } else if (i03 >= args.ne03) { i03 -= args.ne03; }
+
+        int64_t src_idx = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00 + i00;
+        int64_t dst_idx = i3 *args.ne2 *args.ne1 *args.ne0  + i2 *args.ne1 *args.ne0  + i1 *args.ne0  + i0;
+
+        dst_ptr[dst_idx] = src0_ptr[src_idx];
+    }
+}
+
+template <typename T>
+kernel void kernel_pad_impl(
+    constant ggml_metal_kargs_pad & args,
+    device  const char * src0,
+    device        char * dst,
+    uint3 tgpig[[threadgroup_position_in_grid]],
+    uint3 tpitg[[thread_position_in_threadgroup]],
+    uint3   ntg[[threads_per_threadgroup]]) {
+    const int32_t i3 = tgpig.z;
+    const int32_t i2 = tgpig.y;
+    const int32_t k0 = tgpig.x/args.ne1;
+    const int32_t i1 = tgpig.x - k0*args.ne1;
+
+    const int32_t i03 = i3;
+    const int32_t i02 = i2;
+    const int32_t i01 = i1;
+
+    device const T * src0_ptr = (device const T *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
+    device       T * dst_ptr  = (device       T *) (dst  +  i3*args.nb3  +  i2*args.nb2  +  i1*args.nb1);
+
+    for (int32_t l0 = 0; l0 < 1024; l0 += ntg.x) {
+        const int32_t i0 = k0*1024 + tpitg.x + l0;
+        if (i0 >= args.ne0) {
+            break;
+        }
+
+        if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {
+            dst_ptr[i0] = src0_ptr[i0];
+        } else {
+            dst_ptr[i0] = 0.0f;
+        }
+    }
+}
+
+typedef decltype(kernel_pad_impl<float>) kernel_pad_t;
+
+template [[host_name("kernel_pad_f32")]]   kernel kernel_pad_t kernel_pad_impl<float>;
+template [[host_name("kernel_pad_f32_4")]] kernel kernel_pad_t kernel_pad_impl<float4>;
+
+// TODO: this is slow - optimize
+kernel void kernel_pad_reflect_1d_f32(
+    constant   ggml_metal_kargs_pad_reflect_1d & args,
+    device  const char * src0,
+    device        char * dst,
+    uint3 tgpig[[threadgroup_position_in_grid]],
+    uint3  tgpg[[threadgroups_per_grid]],
+    uint3 tpitg[[thread_position_in_threadgroup]],
+    uint3   ntg[[threads_per_threadgroup]]) {
+
+    const int64_t i3 = tgpig.z;
+    const int64_t i2 = tgpig.y;
+    const int64_t i1 = tgpig.x;
+
+    const int64_t i03 = i3;
+    const int64_t i02 = i2;
+    const int64_t i01 = i1;
+
+    device const float * src0_ptr = (device const float *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
+    device       float * dst_ptr  = (device       float *) (dst  +  i3*args.nb3  +  i2*args.nb2  +  i1*args.nb1);
+
+    if (i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {
+        for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
+            if (i0 < args.p0) {
+                dst_ptr[i0] = src0_ptr[args.p0 - i0];
+            } else if (i0 < args.ne0 - args.p1) {
+                dst_ptr[i0] = src0_ptr[i0 - args.p0];
+            } else {
+                dst_ptr[i0] = src0_ptr[(args.ne0 - args.p1 - args.p0) - (args.p1 + 1 - (args.ne0 - i0)) - 1];
+            }
+        }
+    }
+}
+
+kernel void kernel_arange_f32(
+    constant   ggml_metal_kargs_arange & args,
+    device        char * dst,
+    uint3 tgpig[[threadgroup_position_in_grid]],
+    uint3 tpitg[[thread_position_in_threadgroup]],
+    uint3   ntg[[threads_per_threadgroup]]) {
+
+    device float * dst_ptr = (device float *) dst;
+
+    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
+        dst_ptr[i0] = args.start + args.step * i0;
+    }
+}
+
+kernel void kernel_timestep_embedding_f32(
+    constant  ggml_metal_kargs_timestep_embedding & args,
+    device  const char * src0,
+    device        char * dst,
+    uint3 tgpig[[threadgroup_position_in_grid]],
+    uint3 tpitg[[thread_position_in_threadgroup]],
+    uint3   ntg[[threads_per_threadgroup]]) {
+
+    int i = tgpig.x;
+    device float * embed_data = (device float *)(dst + i*args.nb1);
+
+    int half_ = args.dim / 2;
+    for (int j = tpitg.x; j < half_; j += ntg.x) {
+        float timestep = ((device float *)src0)[i];
+        float freq = (float)exp(-log((float)args.max_period) * j / half_);
+        float arg = timestep * freq;
+        embed_data[j        ] = cos(arg);
+        embed_data[j + half_] = sin(arg);
+    }
+
+    if (args.dim % 2 != 0 && tpitg.x == 0) {
+        embed_data[2 * half_] = 0.f;
+    }
+}
+
+kernel void kernel_opt_step_adamw_f32(
+        constant    ggml_metal_kargs_opt_step_adamw & args,
+        device       float * x,
+        device const float * g,
+        device       float * g_m,
+        device       float * g_v,
+        device const float * pars,
+        uint        gid[[thread_position_in_grid]]) {
+
+    if (gid >= args.np) {
+        return;
+    }
+
+    const float alpha  = pars[0];
+    const float beta1  = pars[1];
+    const float beta2  = pars[2];
+    const float eps    = pars[3];
+    const float wd     = pars[4];
+    const float beta1h = pars[5];
+    const float beta2h = pars[6];
+
+    const float gi = g[gid];
+    const float gmi = g_m[gid] * beta1 +      gi * (1.0f - beta1);
+    const float gvi = g_v[gid] * beta2 + gi * gi * (1.0f - beta2);
+
+    g_m[gid] = gmi;
+    g_v[gid] = gvi;
+
+    const float mh =      gmi * beta1h;
+    const float vh = sqrt(gvi * beta2h) + eps;
+
+    x[gid] = x[gid] * (1.0f - alpha * wd) - alpha * mh / vh;
+}
+
+kernel void kernel_opt_step_sgd_f32(
+        constant    ggml_metal_kargs_opt_step_sgd & args,
+        device       float * x,
+        device const float * g,
+        device const float * pars,
+        uint        gid[[thread_position_in_grid]]) {
+
+    if (gid >= args.np) {
+        return;
+    }
+
+    x[gid] = x[gid] * (1.0f - pars[0] * pars[1]) - pars[0] * g[gid];
+}
+
+template<typename T>
+kernel void kernel_memset(
+        constant ggml_metal_kargs_memset & args,
+        device T * dst,
+        uint tpig[[thread_position_in_grid]]) {
+    dst[tpig] = args.val;
+}
+
+typedef decltype(kernel_memset<int64_t>) kernel_memset_t;
+
+template [[host_name("kernel_memset_i64")]] kernel kernel_memset_t kernel_memset<int64_t>;
+
+constant short FC_count_equal_nsg [[function_constant(FC_COUNT_EQUAL + 0)]];
+
+template<typename T>
+kernel void kernel_count_equal(
+        constant ggml_metal_kargs_count_equal & args,
+        device   const char * src0,
+        device   const char * src1,
+        device   atomic_int * dst,
+        threadgroup int32_t * shmem_i32 [[threadgroup(0)]],
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort  tiisg[[thread_index_in_simdgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    const short NSG = FC_count_equal_nsg;
+
+    const int i3 = tgpig.z;
+    const int i2 = tgpig.y;
+    const int i1 = tgpig.x;
+
+    if (i3 >= args.ne03 || i2 >= args.ne02 || i1 >= args.ne01) {
+        return;
+    }
+
+    int sum = 0;
+
+    device const char * base0 = src0 + i1*args.nb01 + i2*args.nb02 + i3*args.nb03;
+    device const char * base1 = src1 + i1*args.nb11 + i2*args.nb12 + i3*args.nb13;
+
+    for (int64_t i0 = tpitg.x; i0 < args.ne00; i0 += ntg.x) {
+        const T v0 = *(device const T *)(base0 + i0*args.nb00);
+        const T v1 = *(device const T *)(base1 + i0*args.nb10);
+        sum += (v0 == v1);
+    }
+
+    sum = simd_sum(sum);
+
+    if (tiisg == 0) {
+        shmem_i32[sgitg] = sum;
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    if (sgitg == 0) {
+        float v = 0.0f;
+        if (tpitg.x < NSG) {
+            v = shmem_i32[tpitg.x];
+        }
+
+        float total = simd_sum(v);
+        if (tpitg.x == 0) {
+            atomic_fetch_add_explicit(dst, (int32_t) total, memory_order_relaxed);
+        }
+    }
+}
+
+typedef decltype(kernel_count_equal<int32_t>) kernel_count_equal_t;
+
+template [[host_name("kernel_count_equal_i32")]] kernel kernel_count_equal_t kernel_count_equal<int32_t>;

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

@@ -0,0 +1,838 @@
+#include "common.h"
+#include "dequantize.h"
+
+constant bool FC_mul_mm_bc_inp [[function_constant(FC_MUL_MM + 0)]];
+constant bool FC_mul_mm_bc_out [[function_constant(FC_MUL_MM + 1)]];
+constant short FC_mul_mm_ne12  [[function_constant(FC_MUL_MM + 2)]];
+constant short FC_mul_mm_ne13  [[function_constant(FC_MUL_MM + 3)]];
+constant short FC_mul_mm_r2    [[function_constant(FC_MUL_MM + 4)]];
+constant short FC_mul_mm_r3    [[function_constant(FC_MUL_MM + 5)]];
+
+// each block_q contains 16*nl weights
+#ifdef GGML_METAL_HAS_TENSOR
+template<
+    typename SA, typename SA_4x4, typename SA_8x8,
+    typename SB, typename SB_2x4, typename SB_8x8,
+    typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread SA_4x4 &),
+    typename T0, typename T0_4x4, typename T1, typename T1_2x4>
+kernel void kernel_mul_mm(
+        constant ggml_metal_kargs_mul_mm & args,
+        device const char * srcA,
+        device const char * srcB,
+        device       char * dst,
+        threadgroup  char * shmem [[threadgroup(0)]],
+        uint3  tgpig [[threadgroup_position_in_grid]],
+        ushort tiitg [[thread_index_in_threadgroup]],
+        ushort sgitg [[simdgroup_index_in_threadgroup]]) {
+    (void) sgitg;
+
+    // Matrix dimensions: A(M,K) x B(K,N) -> C(M,N)
+    const int K = args.ne00;
+    const int M = args.ne0;
+    const int N = args.ne1;
+
+    // Batch dimension handling
+    const int im = tgpig.z;
+    const int i12 = im % FC_mul_mm_ne12;
+    const int i13 = im / FC_mul_mm_ne12;
+
+    // Batch offsets for srcA and srcB
+    const uint64_t offset0 = (i12/FC_mul_mm_r2)*args.nb02 + (i13/FC_mul_mm_r3)*args.nb03;
+
+    // Tile dimensions
+    constexpr int NRB = SZ_SIMDGROUP * N_MM_BLOCK_X * N_MM_SIMD_GROUP_X;
+    constexpr int NRA = SZ_SIMDGROUP * N_MM_BLOCK_Y * N_MM_SIMD_GROUP_Y;
+
+    // Tile offsets in output matrix
+    const int ra = tgpig.y * NRA;
+    const int rb = tgpig.x * NRB;
+
+    // Threadgroup memory for dequantized A tile only
+    threadgroup SA * sa = (threadgroup SA *)(shmem);
+
+    // Work-item count for A loading
+    constexpr int A_WORK_ITEMS = NRA * N_MM_NK;
+    constexpr int NUM_THREADS = N_SIMDWIDTH * N_MM_SIMD_GROUP_X * N_MM_SIMD_GROUP_Y;
+
+    // tA wraps threadgroup memory
+    auto tA = tensor(sa, dextents<int32_t, 2>(N_MM_NK_TOTAL, NRA));
+
+    // tB wraps device memory directly
+    device T1 * ptrB = (device T1 *)(srcB + args.nb12*i12 + args.nb13*i13);
+    const int strideB = args.nb11 / sizeof(T1);
+    auto tB = tensor(ptrB, dextents<int32_t, 2>(K, N), array<int, 2>({1, strideB}));
+
+    // Configure matmul operation
+    mpp::tensor_ops::matmul2d<
+        mpp::tensor_ops::matmul2d_descriptor(
+            NRB, NRA, N_MM_NK_TOTAL, false, true, true,
+            mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
+        execution_simdgroups<N_MM_SIMD_GROUP_X * N_MM_SIMD_GROUP_Y>> mm;
+
+    auto cT = mm.get_destination_cooperative_tensor<decltype(tB), decltype(tA), float>();
+
+    // Accumulate partial results over K dimension
+    for (int loop_k = 0; loop_k < K; loop_k += N_MM_NK_TOTAL) {
+        // === PHASE 1: Dequantization of A into threadgroup memory ===
+        for (int work = tiitg; work < A_WORK_ITEMS; work += NUM_THREADS) {
+            const int row = work / N_MM_NK;
+            const int k_chunk = work % N_MM_NK;
+            const int k_pos = loop_k + k_chunk * 16;
+            const short k_base = k_chunk * 16;
+
+            // Bounds check: skip device read if row is out of matrix bounds
+            if (ra + row < M) {
+                if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
+                    // Element-wise reads when K is not aligned (nb01 not aligned for half4x4/float4x4).
+                    // MSL spec Table 2.5: half4x4 requires 8-byte alignment. When K is odd,
+                    // nb01 = K*2 is not 8-byte aligned, so odd-row pointers are misaligned.
+                    // Mirrors the legacy kernel's existing guard.
+                    device const T0 * row_ptr = (device const T0 *)(srcA + args.nb01 * (ra + row) + offset0);
+
+                    FOR_UNROLL (short i = 0; i < 16; i++) {
+                        sa[row * N_MM_NK_TOTAL + (k_base + i)] = (k_pos + i < K) ? (SA) row_ptr[k_pos + i] : (SA)0;
+                    }
+                } else {
+                    const int block_idx = k_pos / (16 * nl);
+                    const short il = (k_pos / 16) % nl;
+
+                    device const block_q * row_ptr = (device const block_q *)(srcA + args.nb01 * (ra + row) + offset0);
+
+                    SA_4x4 temp_a;
+                    dequantize_func(row_ptr + block_idx, il, temp_a);
+
+                    FOR_UNROLL (short i = 0; i < 16; i++) {
+                        // Zero-pad A for K positions beyond valid range (handles partial K iterations)
+                        sa[row * N_MM_NK_TOTAL + (k_base + i)] = (k_pos + i < K) ? temp_a[i/4][i%4] : (SA)0;
+                    }
+                }
+            } else {
+                // Zero-pad rows beyond matrix bounds
+                FOR_UNROLL (short i = 0; i < 16; i++) {
+                    sa[row * N_MM_NK_TOTAL + (k_base + i)] = (SA)0;
+                }
+            }
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // === PHASE 2: Tensor matmul ===
+        auto mA = tA.slice(0, 0);
+        auto mB = tB.slice(loop_k, rb);
+
+        mm.run(mB, mA, cT);
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    // Store result tile to output matrix (with batch offset)
+    // cT.store handles bounds checking via tD's extents (M, N)
+    device float * dstBatch = (device float *)dst + im * N * M;
+
+    auto tD = tensor(dstBatch, dextents<int32_t, 2>(M, N), array<int, 2>({1, M}));
+    cT.store(tD.slice(ra, rb));
+}
+
+#else
+
+template<
+    typename S0, typename S0_4x4, typename S0_8x8,
+    typename S1, typename S1_2x4, typename S1_8x8,
+    typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread S0_4x4 &),
+    typename T0, typename T0_4x4, typename T1, typename T1_2x4>
+kernel void kernel_mul_mm(
+        constant ggml_metal_kargs_mul_mm & args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        threadgroup  char * shmem [[threadgroup(0)]],
+        uint3  tgpig[[threadgroup_position_in_grid]],
+        ushort tiitg[[thread_index_in_threadgroup]],
+        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
+
+    threadgroup S0 * sa = (threadgroup S0 *)(shmem);
+    threadgroup S1 * sb = (threadgroup S1 *)(shmem + 4096);
+
+    constexpr int NR0 = 64;
+    constexpr int NR1 = 32;
+
+    constexpr int NK  = 32;
+    constexpr int NL0 = NK/16;
+    constexpr int NL1 = NK/8;
+
+    const int im = tgpig.z;
+    const int r0 = tgpig.y*NR0;
+    const int r1 = tgpig.x*NR1;
+
+    // if this block is of 64x32 shape or smaller
+    const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
+    const short nr1 = (args.ne1 - r1 < NR1) ? (args.ne1 - r1) : NR1;
+
+    // a thread shouldn't load data outside of the matrix
+    const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; // 0 .. 63
+    const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; // 0 .. 31
+
+    const short il0 = (tiitg % NL0);
+
+    short il = il0;
+
+    const int i12 = im % FC_mul_mm_ne12;
+    const int i13 = im / FC_mul_mm_ne12;
+
+    const uint64_t offset0 = (i12/FC_mul_mm_r2)*args.nb02 + (i13/FC_mul_mm_r3)*args.nb03;
+    const short    offset1 = il0/nl;
+
+    device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1;
+
+    const short iy = 8*(tiitg % NL1);
+
+    device const T1 * y = (device const T1 *)(src1
+        + args.nb13*i13
+        + args.nb12*i12
+        + args.nb11*(r1 + lr1)
+        + args.nb10*iy);
+
+    S0_8x8 ma[4];
+    S1_8x8 mb[2];
+
+    simdgroup_float8x8 mc[8];
+
+    for (short i = 0; i < 8; i++){
+        mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
+    }
+
+    for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
+        // load data and store to threadgroup memory
+        if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+
+            // no need for dequantization
+            for (short i = 0; i < 16; i++) {
+                const short sx = 2*il0 + i/8;
+                const short sy = (tiitg/NL0)/8;
+
+              //const short lx = i%8;
+              //const short ly = (tiitg/NL0)%8;
+                const short lx = (tiitg/NL0)%8;
+                const short ly = i%8;
+
+                const short ib = 8*sx + sy;
+
+                *(sa + 64*ib + 8*ly + lx) = loop_k + 16*il + i < args.ne00 ? *((device T0 *) x + i) : 0;
+            }
+        } else {
+            S0_4x4 temp_a;
+            dequantize_func(x, il, temp_a);
+
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+
+            FOR_UNROLL (short i = 0; i < 16; i++) {
+                const short sx = 2*il0 + i/8;
+                const short sy = (tiitg/NL0)/8;
+
+              //const short lx = i%8;
+              //const short ly = (tiitg/NL0)%8;
+                const short lx = (tiitg/NL0)%8;
+                const short ly = i%8;
+
+                const short ib = 8*sx + sy;
+
+                // NOTE: this is massively slower.. WTF?
+                //sa[64*ib + 8*ly + lx] = temp_a[i/4][i%4];
+
+                *(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4];
+            }
+        }
+
+        if (FC_mul_mm_bc_inp) {
+            for (short i = 0; i < 8; ++i) {
+                const short sx = (tiitg%NL1);
+                const short sy = (tiitg/NL1)/8;
+
+                const short lx = i;
+                const short ly = (tiitg/NL1)%8;
+              //const short lx = (tiitg/NL1)%8;
+              //const short ly = i;
+
+                const short ib = 4*sx + sy;
+
+                *(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
+            }
+        } else {
+            const short sx = (tiitg%NL1);
+            const short sy = (tiitg/NL1)/8;
+
+          //const short dx = sx;
+          //const short dy = sy;
+
+            const short ly = (tiitg/NL1)%8;
+
+            const short ib = 4*sx + sy;
+
+            *(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)(*((device T1_2x4 *) y));
+        }
+
+        il = (il + 2 < nl) ? il + 2 : il % 2;
+        x  = (il < 2) ? x + (2 + nl - 1)/nl : x;
+
+        y += NK;
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // load matrices from threadgroup memory and conduct outer products
+        threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2));
+        threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2));
+
+        FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
+            simdgroup_barrier(mem_flags::mem_none);
+
+            FOR_UNROLL (short i = 0; i < 4; i++) {
+                simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
+            }
+
+            simdgroup_barrier(mem_flags::mem_none);
+
+            FOR_UNROLL (short i = 0; i < 2; i++) {
+                simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
+            }
+
+            simdgroup_barrier(mem_flags::mem_none);
+
+            FOR_UNROLL (short i = 0; i < 8; i++){
+                simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
+            }
+
+            lsma += 8*64;
+            lsmb += 4*64;
+        }
+    }
+
+    if (!FC_mul_mm_bc_out || (r0 + NR0 <= args.ne0 && r1 + NR1 <= args.ne1)) {
+        // if no bounds checks on the output are needed, we can directly write to device memory
+        device float * C = (device float *) dst +
+            (r0 + 32*(sgitg &  1)) + \
+            (r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0;
+
+        for (short i = 0; i < 8; i++) {
+            simdgroup_store(mc[i], C + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false);
+        }
+    } else {
+        // block is smaller than 64x32, we should avoid writing data outside of the matrix
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
+
+        for (short i = 0; i < 8; i++) {
+            simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (sgitg == 0) {
+            for (int j = tiitg; j < nr1; j += NR1) {
+                device float  * D  = (device float  *) dst + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0;
+                device float4 * D4 = (device float4 *) D;
+
+                threadgroup float  * C  = temp_str + (j*NR0);
+                threadgroup float4 * C4 = (threadgroup float4 *) C;
+
+                int i = 0;
+                for (; i < nr0/4; i++) {
+                    *(D4 + i) = *(C4 + i);
+                }
+
+                i *= 4;
+                for (; i < nr0; i++) {
+                    *(D + i) = *(C + i);
+                }
+            }
+        }
+    }
+}
+
+#endif // GGML_METAL_HAS_TENSOR
+
+template<short ne20> // n_expert_used
+kernel void kernel_mul_mm_id_map0(
+        constant ggml_metal_kargs_mul_mm_id_map0 & args,
+        device  const char * src2,
+        device        char * htpe,
+        device        char * hids,
+        threadgroup   char * shmem [[threadgroup(0)]],
+        ushort tpitg[[thread_position_in_threadgroup]],
+        ushort   ntg[[threads_per_threadgroup]]) {
+    const short ide = tpitg; // expert id
+
+    uint32_t n_all = 0;
+
+    device int32_t * ids_i32 = (device int32_t *) hids + ide*args.ne21;
+
+    for (int i21 = 0; i21 < args.ne21; i21 += ntg) { // n_tokens
+        if (i21 + tpitg < args.ne21) {
+            device const int32_t * src2_i32 = (device const int32_t *) (src2 + (i21 + tpitg)*args.nb21);
+
+            threadgroup uint16_t * sids = (threadgroup uint16_t *) shmem + tpitg*ne20;
+
+            #pragma unroll(ne20)
+            for (short i20 = 0; i20 < ne20; i20++) {
+                sids[i20] = src2_i32[i20];
+            }
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        for (short t = 0; t < ntg; t++) {
+            if (i21 + t >= args.ne21) {
+                break;
+            }
+
+            threadgroup const uint16_t * sids = (threadgroup const uint16_t *) shmem + t*ne20;
+
+            short sel = 0;
+            #pragma unroll(ne20)
+            for (short i20 = 0; i20 < ne20; i20++) {
+                sel += (sids[i20] == ide)*(i20 + 1);
+            }
+
+            ids_i32[n_all] = (i21 + t)*ne20 + sel - 1;
+
+            n_all += sel > 0;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    device uint32_t * tpe_u32 = (device uint32_t *) (htpe);
+    tpe_u32[ide] = n_all;
+}
+
+typedef decltype(kernel_mul_mm_id_map0<1>) kernel_mul_mm_id_map0_t;
+
+template [[host_name("kernel_mul_mm_id_map0_ne20_1" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<1>;
+template [[host_name("kernel_mul_mm_id_map0_ne20_2" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<2>;
+template [[host_name("kernel_mul_mm_id_map0_ne20_4" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<4>;
+template [[host_name("kernel_mul_mm_id_map0_ne20_5" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<5>;
+template [[host_name("kernel_mul_mm_id_map0_ne20_6" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<6>;
+template [[host_name("kernel_mul_mm_id_map0_ne20_8" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<8>;
+template [[host_name("kernel_mul_mm_id_map0_ne20_10")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<10>;
+template [[host_name("kernel_mul_mm_id_map0_ne20_16")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<16>;
+template [[host_name("kernel_mul_mm_id_map0_ne20_22")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<22>;
+
+template<typename S0, typename S0_4x4, typename S0_8x8, typename S1, typename S1_2x4, typename S1_8x8, typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread S0_4x4 &), typename T0, typename T0_4x4, typename T1, typename T1_2x4>
+kernel void kernel_mul_mm_id(
+        constant ggml_metal_kargs_mul_mm_id & args,
+        device const char * src0,
+        device const char * src1,
+        device const char * htpe,
+        device const char * hids,
+        device       char * dst,
+        threadgroup  char * shmem [[threadgroup(0)]],
+        uint3  tgpig[[threadgroup_position_in_grid]],
+        ushort tiitg[[thread_index_in_threadgroup]],
+        ushort tiisg[[thread_index_in_simdgroup]],
+        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
+    threadgroup S0 * sa = (threadgroup S0 *)(shmem);
+    threadgroup S1 * sb = (threadgroup S1 *)(shmem + 4096);
+
+#ifdef GGML_METAL_HAS_TENSOR
+    threadgroup float * sc = (threadgroup float *)(shmem);
+#endif
+
+    constexpr int NR0 = 64;
+    constexpr int NR1 = 32;
+
+    constexpr int NK  = 32;
+    constexpr int NL0 = NK/16;
+    constexpr int NL1 = NK/8;
+
+    const int im = tgpig.z; // expert
+    const int r0 = tgpig.y*NR0;
+    const int r1 = tgpig.x*NR1;
+
+    device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe);
+    device const int32_t  * ids_i32 = (device const int32_t  *) (hids);
+
+    const int32_t neh1 = tpe_u32[im];
+
+    if (r1 >= neh1) {
+        return;
+    }
+
+    // if this block is of 64x32 shape or smaller
+    const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
+    const short nr1 = (    neh1 - r1 < NR1) ? (    neh1 - r1) : NR1;
+
+    // a thread shouldn't load data outside of the matrix
+    const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; // 0 .. 63
+    const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; // 0 .. 31
+
+    const short il0 = (tiitg % NL0);
+
+    short il = il0;
+
+    const int id = ids_i32[im*args.ne21 + r1 + lr1];
+
+    const short i11 = (id % args.ne20) % args.ne11;
+    const short i12 = (id / args.ne20);
+    const short i13 = 0;
+
+    const uint64_t offset0 = im*args.nb02 + i13*args.nb03;
+    const short    offset1 = il0/nl;
+
+    device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1;
+
+    const short iy = 8*(tiitg % NL1);
+
+    device const T1 * y = (device const T1 *)(src1
+        + args.nb13*i13
+        + args.nb12*i12
+        + args.nb11*i11
+        + args.nb10*iy);
+
+#ifndef GGML_METAL_HAS_TENSOR
+    S0_8x8 ma[4];
+    S1_8x8 mb[2];
+
+    simdgroup_float8x8 mc[8];
+
+    for (short i = 0; i < 8; i++){
+        mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
+    }
+#else
+    auto tA = tensor<threadgroup S0, dextents<int32_t, 2>, tensor_inline>(sa, dextents<int32_t, 2>(NK,  NR0));
+    auto tB = tensor<threadgroup S1, dextents<int32_t, 2>, tensor_inline>(sb, dextents<int32_t, 2>(NR1, NK ));
+
+    mpp::tensor_ops::matmul2d<
+        mpp::tensor_ops::matmul2d_descriptor(NR1, NR0, NK, false, true, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
+        execution_simdgroups<4>> mm;
+
+    auto cT = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB), float>();
+#endif
+
+    for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
+#ifndef GGML_METAL_HAS_TENSOR
+        // load data and store to threadgroup memory
+        if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+
+            // no need for dequantization
+            for (short i = 0; i < 16; i++) {
+                const short sx = 2*il0 + i/8;
+                const short sy = (tiitg/NL0)/8;
+
+              //const short lx = i%8;
+              //const short ly = (tiitg/NL0)%8;
+                const short lx = (tiitg/NL0)%8;
+                const short ly = i%8;
+
+                const short ib = 8*sx + sy;
+
+                *(sa + 64*ib + 8*ly + lx) = loop_k + 16*il + i < args.ne00 ? (S0) *((device T0 *) x + i) : (S0) 0;
+            }
+        } else {
+            S0_4x4 temp_a;
+            dequantize_func(x, il, temp_a);
+
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+
+            FOR_UNROLL (short i = 0; i < 16; i++) {
+                const short sx = 2*il0 + i/8;
+                const short sy = (tiitg/NL0)/8;
+
+              //const short lx = i%8;
+              //const short ly = (tiitg/NL0)%8;
+                const short lx = (tiitg/NL0)%8;
+                const short ly = i%8;
+
+                const short ib = 8*sx + sy;
+
+                // NOTE: this is massively slower.. WTF?
+                //sa[64*ib + 8*ly + lx] = temp_a[i/4][i%4];
+
+                *(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4];
+            }
+        }
+
+        if (FC_mul_mm_bc_inp) {
+            for (short i = 0; i < 8; ++i) {
+                const short sx = (tiitg%NL1);
+                const short sy = (tiitg/NL1)/8;
+
+                const short lx = i;
+                const short ly = (tiitg/NL1)%8;
+              //const short lx = (tiitg/NL1)%8;
+              //const short ly = i;
+
+                const short ib = 4*sx + sy;
+
+                *(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
+            }
+        } else {
+            const short sx = (tiitg%NL1);
+            const short sy = (tiitg/NL1)/8;
+
+          //const short dx = sx;
+          //const short dy = sy;
+
+            const short ly = (tiitg/NL1)%8;
+
+            const short ib = 4*sx + sy;
+
+            *(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)(*((device T1_2x4 *) y));
+        }
+#else
+        // load data and store to threadgroup memory
+        if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+
+            // no need for dequantization
+            for (short i = 0; i < 16; i++) {
+                const short sx = 2*il0 + i/8;
+                const short sy = (tiitg/NL0)/8;
+
+                const short lx = i%8;
+                const short ly = (tiitg/NL0)%8;
+                //const short lx = (tiitg/NL0)%8;
+                //const short ly = i%8;
+
+                *(sa + NK*(8*sy + ly) + 8*sx + lx) = loop_k + 16*il + i < args.ne00 ? *((device T0 *) x + i) : 0;
+            }
+        } else {
+            S0_4x4 temp_a;
+            dequantize_func(x, il, temp_a);
+
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+
+            FOR_UNROLL (short i = 0; i < 16; i++) {
+                const short sx = 2*il0 + i/8;
+                const short sy = (tiitg/NL0)/8;
+
+                const short lx = i%8;
+                const short ly = (tiitg/NL0)%8;
+                //const short lx = (tiitg/NL0)%8;
+                //const short ly = i%8;
+
+                *(sa + NK*(8*sy + ly) + 8*sx + lx) = temp_a[i/4][i%4];
+            }
+        }
+
+        if (FC_mul_mm_bc_inp) {
+            for (short i = 0; i < 8; ++i) {
+                const short sx = (tiitg%NL1);
+                const short sy = (tiitg/NL1)/8;
+
+                const short lx = i;
+                const short ly = (tiitg/NL1)%8;
+                //const short lx = (tiitg/NL1)%8;
+                //const short ly = i;
+
+                *(sb + NK*(8*sy + ly) + 8*sx + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
+            }
+        } else {
+            const short sx = (tiitg%NL1);
+            const short sy = (tiitg/NL1)/8;
+
+            //const short lx = i;
+            const short ly = (tiitg/NL1)%8;
+            //const short lx = (tiitg/NL1)%8;
+            //const short ly = i;
+
+            *(threadgroup S1_2x4 *)(sb + NK*(8*sy + ly) + 8*sx) = (S1_2x4)(*((device T1_2x4 *) y));
+        }
+#endif
+
+        il = (il + 2 < nl) ? il + 2 : il % 2;
+        x  = (il < 2) ? x + (2 + nl - 1)/nl : x;
+
+        y += NK;
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+#ifndef GGML_METAL_HAS_TENSOR
+        // load matrices from threadgroup memory and conduct outer products
+        threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2));
+        threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2));
+
+        FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
+            simdgroup_barrier(mem_flags::mem_none);
+
+            FOR_UNROLL (short i = 0; i < 4; i++) {
+                simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
+            }
+
+            simdgroup_barrier(mem_flags::mem_none);
+
+            FOR_UNROLL (short i = 0; i < 2; i++) {
+                simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
+            }
+
+            simdgroup_barrier(mem_flags::mem_none);
+
+            FOR_UNROLL (short i = 0; i < 8; i++){
+                simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
+            }
+
+            lsma += 8*64;
+            lsmb += 4*64;
+        }
+#else
+        auto sA = tA.slice(0, 0);
+        auto sB = tB.slice(0, 0);
+
+        mm.run(sB, sA, cT);
+#endif
+    }
+
+    // block is smaller than 64x32, we should avoid writing data outside of the matrix
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+#ifdef GGML_METAL_HAS_TENSOR
+    auto tC = tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline>(sc, dextents<int32_t, 2>(NR0, NR1));
+    cT.store(tC);
+#else
+    threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
+
+    for (short i = 0; i < 8; i++) {
+        simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
+    }
+#endif
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    for (short j = sgitg; j < nr1; j += 4) {
+        const int id = ids_i32[im*args.ne21 + r1 + j];
+
+        const short ide = id % args.ne20;
+        const short idt = id / args.ne20;
+
+        device float  * D  = (device float  *) dst + r0 + ide*args.ne0 + idt*args.ne1*args.ne0;
+        device float4 * D4 = (device float4 *) D;
+
+        threadgroup float  * C  = (threadgroup float  *) shmem + j*NR0;
+        threadgroup float4 * C4 = (threadgroup float4 *) C;
+
+        int i = tiisg;
+        for (; i < nr0/4; i += 32) {
+            *(D4 + i) = *(C4 + i);
+        }
+
+        i = (4*(nr0/4)) + tiisg;
+        for (; i < nr0; i += 32) {
+            *(D + i) = *(C + i);
+        }
+    }
+}
+
+//
+// matrix-matrix multiplication
+//
+
+typedef decltype(kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, float4x4, 1, dequantize_f32, float, float4x4, float, float2x4>) mul_mm_t;
+
+template [[host_name("kernel_mul_mm_f32_f32")]]     kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   float4x4,      1,     dequantize_f32,     float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_f16_f32")]]     kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   half4x4,       1,     dequantize_f16,     half,   half4x4,   float, float2x4>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_mul_mm_bf16_f32")]]    kernel mul_mm_t kernel_mul_mm<bfloat, bfloat4x4, simdgroup_bfloat8x8, bfloat, bfloat2x4, simdgroup_bfloat8x8, bfloat4x4,     1,     dequantize_bf16,    bfloat, bfloat4x4, float, float2x4>;
+#endif
+template [[host_name("kernel_mul_mm_q1_0_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q1_0,    8,     dequantize_q1_0,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_q4_0_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_0,    2,     dequantize_q4_0,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_q4_1_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_1,    2,     dequantize_q4_1,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_q5_0_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_0,    2,     dequantize_q5_0,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_q5_1_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_1,    2,     dequantize_q5_1,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_q8_0_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q8_0,    2,     dequantize_q8_0,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_mxfp4_f32")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_mxfp4,   2,     dequantize_mxfp4,   float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_q2_K_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q2_K,    QK_NL, dequantize_q2_K,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_q3_K_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q3_K,    QK_NL, dequantize_q3_K,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_q4_K_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_K,    QK_NL, dequantize_q4_K,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_q5_K_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_K,    QK_NL, dequantize_q5_K,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_q6_K_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q6_K,    QK_NL, dequantize_q6_K,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_iq2_xxs_f32")]] kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_iq2_xs_f32")]]  kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_xs,  QK_NL, dequantize_iq2_xs,  float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_iq3_xxs_f32")]] kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq3_xxs, QK_NL, dequantize_iq3_xxs, float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_iq3_s_f32")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq3_s,   QK_NL, dequantize_iq3_s,   float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_iq2_s_f32")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_s,   QK_NL, dequantize_iq2_s,   float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_iq1_s_f32")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq1_s,   QK_NL, dequantize_iq1_s,   float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_iq1_m_f32")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq1_m,   QK_NL, dequantize_iq1_m,   float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_iq4_nl_f32")]]  kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq4_nl,  2,     dequantize_iq4_nl,  float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_iq4_xs_f32")]]  kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq4_xs,  QK_NL, dequantize_iq4_xs,  float,  float4x4,  float, float2x4>;
+
+template [[host_name("kernel_mul_mm_f32_f16")]]     kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   float4x4,      1,     dequantize_f32,     float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_f16_f16")]]     kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   half4x4,       1,     dequantize_f16,     half,   half4x4,   half, half2x4>;
+template [[host_name("kernel_mul_mm_q1_0_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q1_0,    8,     dequantize_q1_0,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_q4_0_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_0,    2,     dequantize_q4_0,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_q4_1_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_1,    2,     dequantize_q4_1,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_q5_0_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_0,    2,     dequantize_q5_0,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_q5_1_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_1,    2,     dequantize_q5_1,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_q8_0_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q8_0,    2,     dequantize_q8_0,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_mxfp4_f16")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_mxfp4,   2,     dequantize_mxfp4,   float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_q2_K_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q2_K,    QK_NL, dequantize_q2_K,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_q3_K_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q3_K,    QK_NL, dequantize_q3_K,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_q4_K_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_K,    QK_NL, dequantize_q4_K,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_q5_K_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_K,    QK_NL, dequantize_q5_K,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_q6_K_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q6_K,    QK_NL, dequantize_q6_K,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_iq2_xxs_f16")]] kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_iq2_xs_f16")]]  kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_xs,  QK_NL, dequantize_iq2_xs,  float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_iq3_xxs_f16")]] kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq3_xxs, QK_NL, dequantize_iq3_xxs, float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_iq3_s_f16")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq3_s,   QK_NL, dequantize_iq3_s,   float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_iq2_s_f16")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_s,   QK_NL, dequantize_iq2_s,   float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_iq1_s_f16")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq1_s,   QK_NL, dequantize_iq1_s,   float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_iq1_m_f16")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq1_m,   QK_NL, dequantize_iq1_m,   float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_iq4_nl_f16")]]  kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq4_nl,  2,     dequantize_iq4_nl,  float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_iq4_xs_f16")]]  kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq4_xs,  QK_NL, dequantize_iq4_xs,  float,  float4x4,  half, half2x4>;
+
+//
+// indirect matrix-matrix multiplication
+//
+
+typedef decltype(kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, float4x4, 1, dequantize_f32, float, float4x4, float, float2x4>) mul_mm_id;
+
+template [[host_name("kernel_mul_mm_id_f32_f32")]]     kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   float4x4,      1,     dequantize_f32,     float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_f16_f32")]]     kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   half4x4,       1,     dequantize_f16,     half,   half4x4,   float, float2x4>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_mul_mm_id_bf16_f32")]]    kernel mul_mm_id kernel_mul_mm_id<bfloat, bfloat4x4, simdgroup_bfloat8x8, bfloat, bfloat2x4, simdgroup_bfloat8x8, bfloat4x4,     1,     dequantize_bf16,    bfloat, bfloat4x4, float, float2x4>;
+#endif
+template [[host_name("kernel_mul_mm_id_q1_0_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q1_0,    8,     dequantize_q1_0,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_q4_0_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_0,    2,     dequantize_q4_0,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_q4_1_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_1,    2,     dequantize_q4_1,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_q5_0_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_0,    2,     dequantize_q5_0,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_q5_1_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_1,    2,     dequantize_q5_1,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_q8_0_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q8_0,    2,     dequantize_q8_0,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_mxfp4_f32")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_mxfp4,   2,     dequantize_mxfp4,   float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_q2_K_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q2_K,    QK_NL, dequantize_q2_K,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_q3_K_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q3_K,    QK_NL, dequantize_q3_K,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_q4_K_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_K,    QK_NL, dequantize_q4_K,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_q5_K_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_K,    QK_NL, dequantize_q5_K,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_q6_K_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q6_K,    QK_NL, dequantize_q6_K,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_iq2_xxs_f32")]] kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_iq2_xs_f32")]]  kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_xs,  QK_NL, dequantize_iq2_xs,  float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_iq3_xxs_f32")]] kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq3_xxs, QK_NL, dequantize_iq3_xxs, float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_iq3_s_f32")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq3_s,   QK_NL, dequantize_iq3_s,   float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_iq2_s_f32")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_s,   QK_NL, dequantize_iq2_s,   float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_iq1_s_f32")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq1_s,   QK_NL, dequantize_iq1_s,   float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_iq1_m_f32")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq1_m,   QK_NL, dequantize_iq1_m,   float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_iq4_nl_f32")]]  kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq4_nl,  2,     dequantize_iq4_nl,  float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_iq4_xs_f32")]]  kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq4_xs,  QK_NL, dequantize_iq4_xs,  float,  float4x4,  float, float2x4>;
+
+template [[host_name("kernel_mul_mm_id_f32_f16")]]     kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   float4x4,      1,     dequantize_f32,     float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_f16_f16")]]     kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   half4x4,       1,     dequantize_f16,     half,   half4x4,   half, half2x4>;
+template [[host_name("kernel_mul_mm_id_q1_0_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q1_0,    8,     dequantize_q1_0,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_q4_0_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_0,    2,     dequantize_q4_0,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_q4_1_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_1,    2,     dequantize_q4_1,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_q5_0_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_0,    2,     dequantize_q5_0,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_q5_1_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_1,    2,     dequantize_q5_1,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_q8_0_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q8_0,    2,     dequantize_q8_0,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_mxfp4_f16")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_mxfp4,   2,     dequantize_mxfp4,   float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_q2_K_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q2_K,    QK_NL, dequantize_q2_K,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_q3_K_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q3_K,    QK_NL, dequantize_q3_K,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_q4_K_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_K,    QK_NL, dequantize_q4_K,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_q5_K_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_K,    QK_NL, dequantize_q5_K,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_q6_K_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q6_K,    QK_NL, dequantize_q6_K,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_iq2_xxs_f16")]] kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_iq2_xs_f16")]]  kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_xs,  QK_NL, dequantize_iq2_xs,  float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_iq3_xxs_f16")]] kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq3_xxs, QK_NL, dequantize_iq3_xxs, float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_iq3_s_f16")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq3_s,   QK_NL, dequantize_iq3_s,   float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_iq2_s_f16")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq2_s,   QK_NL, dequantize_iq2_s,   float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_iq1_s_f16")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq1_s,   QK_NL, dequantize_iq1_s,   float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_iq1_m_f16")]]   kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq1_m,   QK_NL, dequantize_iq1_m,   float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_iq4_nl_f16")]]  kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq4_nl,  2,     dequantize_iq4_nl,  float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_iq4_xs_f16")]]  kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq4_xs,  QK_NL, dequantize_iq4_xs,  float,  float4x4,  half, half2x4>;

ggml/src/ggml-metal/kernels/norm.metal

@@ -0,0 +1,308 @@
+#include "common.h"
+
+// F == 1 : norm (no fuse)
+// F == 2 : norm + mul
+// F == 3 : norm + mul + add
+template <typename T, short F>
+kernel void kernel_norm_fuse_impl(
+        constant ggml_metal_kargs_norm & args,
+        device const char * src0,
+        device const char * src1_0,
+        device const char * src1_1,
+        device       char * dst,
+        threadgroup float * shmem_f32 [[threadgroup(0)]],
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort  tiisg[[thread_index_in_simdgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    if (sgitg == 0) {
+        shmem_f32[tiisg] = 0.0f;
+    }
+
+    const int i01 = tgpig.x;
+    const int i02 = tgpig.y;
+    const int i03 = tgpig.z;
+
+    device const T * x = (device const T *) (src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
+
+    device const T * f0 = (device const T *) (src1_0 + (i03%args.nef3[1])*args.nbf3[1] + (i02%args.nef2[1])*args.nbf2[1] + (i01%args.nef1[1])*args.nbf1[1]);
+    device const T * f1 = (device const T *) (src1_1 + (i03%args.nef3[2])*args.nbf3[2] + (i02%args.nef2[2])*args.nbf2[2] + (i01%args.nef1[2])*args.nbf1[2]);
+
+    T sumft(0.0f);
+
+    float sumf = 0.0f;
+
+    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
+        sumft += x[i00];
+    }
+    sumf = dot(sumft, T(1.0f));
+    sumf = simd_sum(sumf);
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    if (tiisg == 0) {
+        shmem_f32[sgitg] = sumf;
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    sumf = shmem_f32[tiisg];
+    sumf = simd_sum(sumf);
+
+    const float mean = sumf/args.ne00;
+
+    device T * y = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);
+
+    sumf = 0.0f;
+    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
+        y[i00] = x[i00] - mean;
+        sumf += dot(y[i00], y[i00]);
+    }
+    sumf = simd_sum(sumf);
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    if (tiisg == 0) {
+        shmem_f32[sgitg] = sumf;
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    sumf = shmem_f32[tiisg];
+    sumf = simd_sum(sumf);
+
+    const float variance = sumf/args.ne00;
+
+    const float scale = 1.0f/sqrt(variance + args.eps);
+    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
+        if (F == 1) {
+            y[i00] = (y[i00]*scale);
+        }
+        if (F == 2) {
+            y[i00] = (y[i00]*scale)*f0[i00];
+        }
+        if (F == 3) {
+            y[i00] = (y[i00]*scale)*f0[i00] + f1[i00];
+        }
+    }
+}
+
+typedef decltype(kernel_norm_fuse_impl<float4, 1>) kernel_norm_fuse_t;
+
+template [[host_name("kernel_norm_f32")]]         kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float, 1>;
+template [[host_name("kernel_norm_mul_f32")]]     kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float, 2>;
+template [[host_name("kernel_norm_mul_add_f32")]] kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float, 3>;
+
+template [[host_name("kernel_norm_f32_4")]]         kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float4, 1>;
+template [[host_name("kernel_norm_mul_f32_4")]]     kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float4, 2>;
+template [[host_name("kernel_norm_mul_add_f32_4")]] kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float4, 3>;
+
+// F == 1 : rms_norm (no fuse)
+// F == 2 : rms_norm + mul
+// F == 3 : rms_norm + mul + add
+template <typename T, short F>
+kernel void kernel_rms_norm_fuse_impl(
+        constant ggml_metal_kargs_norm & args,
+        device const char * src0,
+        device const char * src1_0,
+        device const char * src1_1,
+        device       char * dst,
+        threadgroup float * shmem_f32 [[threadgroup(0)]],
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort  tiisg[[thread_index_in_simdgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    if (sgitg == 0) {
+        shmem_f32[tiisg] = 0.0f;
+    }
+
+    const int i01 = tgpig.x;
+    const int i02 = tgpig.y;
+    const int i03 = tgpig.z;
+
+    device const T * x = (device const T *) (src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
+
+    device const T * f0 = (device const T *) (src1_0 + (i03%args.nef3[1])*args.nbf3[1] + (i02%args.nef2[1])*args.nbf2[1] + (i01%args.nef1[1])*args.nbf1[1]);
+    device const T * f1 = (device const T *) (src1_1 + (i03%args.nef3[2])*args.nbf3[2] + (i02%args.nef2[2])*args.nbf2[2] + (i01%args.nef1[2])*args.nbf1[2]);
+
+    float sumf = 0.0f;
+
+    // parallel sum
+    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
+        sumf += dot(x[i00], x[i00]);
+    }
+    sumf = simd_sum(sumf);
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    if (tiisg == 0) {
+        shmem_f32[sgitg] = sumf;
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    sumf = shmem_f32[tiisg];
+    sumf = simd_sum(sumf);
+
+    const float mean  = sumf/args.ne00;
+    const float scale = 1.0f/sqrt(mean + args.eps);
+
+    device T * y = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);
+    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
+        if (F == 1) {
+            y[i00] = (x[i00]*scale);
+        }
+        if (F == 2) {
+            y[i00] = (x[i00]*scale)*f0[i00];
+        }
+        if (F == 3) {
+            y[i00] = (x[i00]*scale)*f0[i00] + f1[i00];
+        }
+    }
+}
+
+typedef decltype(kernel_rms_norm_fuse_impl<float4, 1>) kernel_rms_norm_fuse_t;
+
+template [[host_name("kernel_rms_norm_f32")]]         kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float, 1>;
+template [[host_name("kernel_rms_norm_mul_f32")]]     kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float, 2>;
+template [[host_name("kernel_rms_norm_mul_add_f32")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float, 3>;
+
+template [[host_name("kernel_rms_norm_f32_4")]]         kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 1>;
+template [[host_name("kernel_rms_norm_mul_f32_4")]]     kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 2>;
+template [[host_name("kernel_rms_norm_mul_add_f32_4")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 3>;
+
+template <typename T0, typename T>
+kernel void kernel_l2_norm_impl(
+        constant ggml_metal_kargs_l2_norm & args,
+        device const char * src0,
+        device       char * dst,
+        threadgroup float * shmem_f32 [[threadgroup(0)]],
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort  tiisg[[thread_index_in_simdgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    const int i03 = tgpig.z;
+    const int i02 = tgpig.y;
+    const int i01 = tgpig.x;
+
+    if (sgitg == 0) {
+        shmem_f32[tiisg] = 0.0f;
+    }
+
+    device const T0 * x = (device const T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
+    device       T  * y = (device       T  *) (dst  + i03*args.nb3  + i02*args.nb2  + i01*args.nb1);
+
+    float sumf = 0.0f;
+
+    // parallel sum
+    for (int i00 = tpitg.x; i00 < args.ne00; i00 += ntg.x) {
+        sumf += dot(x[i00], x[i00]);
+    }
+    sumf = simd_sum(sumf);
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    if (tiisg == 0) {
+        shmem_f32[sgitg] = sumf;
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    sumf = shmem_f32[tiisg];
+    sumf = simd_sum(sumf);
+
+    const float scale = 1.0f/max(sqrt(sumf), args.eps);
+
+    for (int i00 = tpitg.x; i00 < args.ne00; i00 += ntg.x) {
+        y[i00] = x[i00] * scale;
+    }
+}
+
+typedef decltype(kernel_l2_norm_impl<float, float>) kernel_l2_norm_t;
+
+template [[host_name("kernel_l2_norm_f32_f32")]]   kernel kernel_l2_norm_t kernel_l2_norm_impl<float,  float>;
+template [[host_name("kernel_l2_norm_f32_f32_4")]] kernel kernel_l2_norm_t kernel_l2_norm_impl<float4, float4>;
+
+kernel void kernel_group_norm_f32(
+        constant ggml_metal_kargs_group_norm & args,
+        device const float * src0,
+        device       float * dst,
+        threadgroup float  * buf [[threadgroup(0)]],
+        uint tgpig[[threadgroup_position_in_grid]],
+        uint tpitg[[thread_position_in_threadgroup]],
+        uint sgitg[[simdgroup_index_in_threadgroup]],
+        uint tiisg[[thread_index_in_simdgroup]],
+        uint   ntg[[threads_per_threadgroup]]) {
+    const int64_t ne = args.ne00*args.ne01*args.ne02;
+    const int64_t gs = args.ne00*args.ne01*((args.ne02 + args.ngrp - 1) / args.ngrp);
+
+    int start = tgpig * gs;
+    int end   = start + gs;
+
+    start += tpitg;
+
+    if (end >= ne) {
+        end = ne;
+    }
+
+    float tmp = 0.0f; // partial sum for thread in warp
+
+    for (int j = start; j < end; j += ntg) {
+        tmp += src0[j];
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    tmp = simd_sum(tmp);
+    if (ntg > N_SIMDWIDTH) {
+        if (sgitg == 0) {
+            buf[tiisg] = 0.0f;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiisg == 0) {
+            buf[sgitg] = tmp;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        tmp = buf[tiisg];
+        tmp = simd_sum(tmp);
+    }
+
+    const float mean = tmp / gs;
+    tmp = 0.0f;
+
+    for (int j = start; j < end; j += ntg) {
+        float xi = src0[j] - mean;
+        dst[j] = xi;
+        tmp += xi * xi;
+    }
+
+    tmp = simd_sum(tmp);
+    if (ntg > N_SIMDWIDTH) {
+        if (sgitg == 0) {
+            buf[tiisg] = 0.0f;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiisg == 0) {
+            buf[sgitg] = tmp;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        tmp = buf[tiisg];
+        tmp = simd_sum(tmp);
+    }
+
+    const float variance = tmp / gs;
+    const float scale = 1.0f/sqrt(variance + args.eps);
+    for (int j = start; j < end; j += ntg) {
+        dst[j] *= scale;
+    }
+}

ggml/src/ggml-metal/kernels/pool.metal

@@ -0,0 +1,148 @@
+#include "common.h"
+
+kernel void kernel_pool_2d_max_f32(
+        constant    ggml_metal_kargs_pool_2d & args,
+        device  const float * src0,
+        device        float * dst,
+        uint        gid[[thread_position_in_grid]]) {
+
+    if (gid >= args.np) {
+        return;
+    }
+
+    const int idx = gid;
+    const int I_HW = args.IH * args.IW;
+    const int O_HW = args.OH * args.OW;
+    const int nc = idx / O_HW;
+    const int cur_oh = idx % O_HW / args.OW;
+    const int cur_ow = idx % O_HW % args.OW;
+
+    device const float * i_ptr = src0 + nc * I_HW;
+    device       float * o_ptr = dst  + nc * O_HW;
+
+    const int start_h = cur_oh * args.s1 - args.p1;
+    const int bh = MAX(0,  start_h);
+    const int eh = MIN(args.IH, start_h + args.k1);
+    const int start_w = cur_ow * args.s0 - args.p0;
+    const int bw = MAX(0,  start_w);
+    const int ew = MIN(args.IW, start_w + args.k0);
+
+    float res = -INFINITY;
+
+    for (int i = bh; i < eh; i += 1) {
+        for (int j = bw; j < ew; j += 1) {
+            res = MAX(res, i_ptr[i * args.IW + j]);
+        }
+    }
+
+    o_ptr[cur_oh * args.OW + cur_ow] = res;
+}
+
+kernel void kernel_pool_2d_avg_f32(
+        constant    ggml_metal_kargs_pool_2d & args,
+        device  const float * src0,
+        device        float * dst,
+        uint        gid[[thread_position_in_grid]]) {
+
+    if (gid >= args.np) {
+        return;
+    }
+
+    const int idx = gid;
+    const int I_HW = args.IH * args.IW;
+    const int O_HW = args.OH * args.OW;
+    const int nc = idx / O_HW;
+    const int cur_oh = idx % O_HW / args.OW;
+    const int cur_ow = idx % O_HW % args.OW;
+
+    device const float * i_ptr = src0 + nc * I_HW;
+    device       float * o_ptr = dst  + nc * O_HW;
+
+    const int start_h = cur_oh * args.s1 - args.p1;
+    const int bh = MAX(0,  start_h);
+    const int eh = MIN(args.IH, start_h + args.k1);
+    const int start_w = cur_ow * args.s0 - args.p0;
+    const int bw = MAX(0,  start_w);
+    const int ew = MIN(args.IW, start_w + args.k0);
+    // const float scale = 1. / ((eh - bh) * (ew - bw));
+    const float scale = 1. / (args.k0 * args.k1);
+
+    float res = 0;
+
+    for (int i = bh; i < eh; i += 1) {
+        for (int j = bw; j < ew; j += 1) {
+            float cur = i_ptr[i * args.IW + j];
+            res += cur * scale;
+        }
+    }
+
+    o_ptr[cur_oh * args.OW + cur_ow] = res;
+}
+
+
+kernel void kernel_pool_1d_max_f32(
+        constant        ggml_metal_kargs_pool_1d & args,
+        device  const   float * src,
+        device          float * dst,
+        uint            gid [[thread_position_in_grid]]
+) {
+
+    if (gid >= args.np) {
+        return;
+    }
+
+    const int ow  = (int)gid % args.OW;
+    const int row = (int)gid / args.OW;
+
+    const int base = ow * args.s0 - args.p0;
+
+    float acc = -INFINITY;
+
+    const int src_off = row * args.IW;
+    const int dst_off = row * args.OW;
+
+    for (int ki = 0; ki < args.k0; ++ki) {
+        int j = base + ki;
+        if (j < 0 || j >= args.IW){
+            continue;
+        }
+        float v = src[src_off + j];
+        acc = max(acc, v);
+    }
+
+    dst[dst_off + ow] = acc;
+}
+
+kernel void kernel_pool_1d_avg_f32(
+        constant        ggml_metal_kargs_pool_1d & args,
+        device  const   float * src,
+        device          float * dst,
+        uint            gid [[thread_position_in_grid]]
+) {
+
+    if (gid >= args.np) {
+        return;
+    }
+
+    const int ow  = (int)gid % args.OW;
+    const int row = (int)gid / args.OW;
+
+    const int base = ow * args.s0 - args.p0;
+
+    float acc = 0.0f;
+    int   cnt = 0;
+
+    const int src_off = row * args.IW;
+    const int dst_off = row * args.OW;
+
+    for (int ki = 0; ki < args.k0; ++ki) {
+        const int j = base + ki;
+        if (j < 0 || j >= args.IW) {
+            continue;
+        }
+        acc += src[src_off + j];
+        cnt += 1;
+    }
+
+    dst[dst_off + ow] = (cnt > 0) ? (acc / (float)cnt) : 0.0f;
+}

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

@@ -0,0 +1,213 @@
+#pragma once
+
+#include "common.h"
+
+void quantize_q1_0(device const float * src, device block_q1_0 & dst) {
+    float sum_abs = 0.0f;
+    for (int j = 0; j < QK1_0; j++) {
+        sum_abs += fabs(src[j]);
+    }
+    dst.d = sum_abs / QK1_0;
+
+    for (int j = 0; j < QK1_0 / 8; j++) {
+        dst.qs[j] = 0;
+    }
+    for (int j = 0; j < QK1_0; j++) {
+        if (src[j] >= 0.0f) {
+            dst.qs[j / 8] |= (1 << (j % 8));
+        }
+    }
+}
+
+void quantize_q4_0(device const float * src, device block_q4_0 & dst) {
+#pragma METAL fp math_mode(safe)
+    float amax = 0.0f; // absolute max
+    float max  = 0.0f;
+
+    for (int j = 0; j < QK4_0; j++) {
+        const float v = src[j];
+        if (amax < fabs(v)) {
+            amax = fabs(v);
+            max  = v;
+        }
+    }
+
+    const float d = max / -8;
+    const float id = d ? 1.0f/d : 0.0f;
+
+    dst.d = d;
+
+    for (int j = 0; j < QK4_0/2; ++j) {
+        const float x0 = src[0       + j]*id;
+        const float x1 = src[QK4_0/2 + j]*id;
+
+        const uint8_t xi0 = MIN(15, (int8_t)(x0 + 8.5f));
+        const uint8_t xi1 = MIN(15, (int8_t)(x1 + 8.5f));
+
+        dst.qs[j]  = xi0;
+        dst.qs[j] |= xi1 << 4;
+    }
+}
+
+void quantize_q4_1(device const float * src, device block_q4_1 & dst) {
+#pragma METAL fp math_mode(safe)
+    float min = FLT_MAX;
+    float max = -FLT_MAX;
+
+    for (int j = 0; j < QK4_1; j++) {
+        const float v = src[j];
+        if (min > v) min = v;
+        if (max < v) max = v;
+    }
+
+    const float d = (max - min) / ((1 << 4) - 1);
+    const float id = d ? 1.0f/d : 0.0f;
+
+    dst.d = d;
+    dst.m = min;
+
+    for (int j = 0; j < QK4_1/2; ++j) {
+        const float x0 = (src[0       + j] - min)*id;
+        const float x1 = (src[QK4_1/2 + j] - min)*id;
+
+        const uint8_t xi0 = MIN(15, (int8_t)(x0 + 0.5f));
+        const uint8_t xi1 = MIN(15, (int8_t)(x1 + 0.5f));
+
+        dst.qs[j]  = xi0;
+        dst.qs[j] |= xi1 << 4;
+    }
+}
+
+void quantize_q5_0(device const float * src, device block_q5_0 & dst) {
+#pragma METAL fp math_mode(safe)
+    float amax = 0.0f; // absolute max
+    float max  = 0.0f;
+
+    for (int j = 0; j < QK5_0; j++) {
+        const float v = src[j];
+        if (amax < fabs(v)) {
+            amax = fabs(v);
+            max  = v;
+        }
+    }
+
+    const float d = max / -16;
+    const float id = d ? 1.0f/d : 0.0f;
+
+    dst.d = d;
+
+    uint32_t qh = 0;
+    for (int j = 0; j < QK5_0/2; ++j) {
+        const float x0 = src[0       + j]*id;
+        const float x1 = src[QK5_0/2 + j]*id;
+
+        const uint8_t xi0 = MIN(31, (int8_t)(x0 + 16.5f));
+        const uint8_t xi1 = MIN(31, (int8_t)(x1 + 16.5f));
+
+        dst.qs[j] = (xi0 & 0xf) | ((xi1 & 0xf) << 4);
+        qh |= ((xi0 & 0x10u) >> 4) << (j + 0);
+        qh |= ((xi1 & 0x10u) >> 4) << (j + QK5_0/2);
+    }
+
+    thread const uint8_t * qh8 = (thread const uint8_t *)&qh;
+
+    for (int j = 0; j < 4; ++j) {
+        dst.qh[j] = qh8[j];
+    }
+}
+
+void quantize_q5_1(device const float * src, device block_q5_1 & dst) {
+#pragma METAL fp math_mode(safe)
+    float max = src[0];
+    float min = src[0];
+
+    for (int j = 1; j < QK5_1; j++) {
+        const float v = src[j];
+        min = v < min ? v : min;
+        max = v > max ? v : max;
+    }
+
+    const float d = (max - min) / 31;
+    const float id = d ? 1.0f/d : 0.0f;
+
+    dst.d = d;
+    dst.m = min;
+
+    uint32_t qh = 0;
+    for (int j = 0; j < QK5_1/2; ++j) {
+        const float x0 = (src[0       + j] - min)*id;
+        const float x1 = (src[QK5_1/2 + j] - min)*id;
+
+        const uint8_t xi0 = (uint8_t)(x0 + 0.5f);
+        const uint8_t xi1 = (uint8_t)(x1 + 0.5f);
+
+        dst.qs[j] = (xi0 & 0xf) | ((xi1 & 0xf) << 4);
+        qh |= ((xi0 & 0x10u) >> 4) << (j + 0);
+        qh |= ((xi1 & 0x10u) >> 4) << (j + QK5_1/2);
+    }
+
+    thread const uint8_t * qh8 = (thread const uint8_t *)&qh;
+
+    for (int j = 0; j < 4; ++j) {
+        dst.qh[j] = qh8[j];
+    }
+}
+
+void quantize_q8_0(device const float * src, device block_q8_0 & dst) {
+#pragma METAL fp math_mode(safe)
+    float amax = 0.0f; // absolute max
+
+    for (int j = 0; j < QK8_0; j++) {
+        const float v = src[j];
+        amax = MAX(amax, fabs(v));
+    }
+
+    const float d = amax / ((1 << 7) - 1);
+    const float id = d ? 1.0f/d : 0.0f;
+
+    dst.d = d;
+
+    for (int j = 0; j < QK8_0; ++j) {
+        const float x0 = src[j]*id;
+
+        dst.qs[j] = round(x0);
+    }
+}
+
+void quantize_iq4_nl(device const float * src, device block_iq4_nl & dst) {
+#pragma METAL fp math_mode(safe)
+    float amax = 0.0f; // absolute max
+    float max  = 0.0f;
+
+    for (int j = 0; j < QK4_NL; j++) {
+        const float v = src[j];
+        if (amax < fabs(v)) {
+            amax = fabs(v);
+            max  = v;
+        }
+    }
+
+    const float d = max / kvalues_iq4nl_f[0];
+    const float id = d ? 1.0f/d : 0.0f;
+
+    float sumqx = 0, sumq2 = 0;
+    for (int j = 0; j < QK4_NL/2; ++j) {
+        const float x0 = src[0        + j]*id;
+        const float x1 = src[QK4_NL/2 + j]*id;
+
+        const uint8_t xi0 = best_index_int8(16, kvalues_iq4nl_f, x0);
+        const uint8_t xi1 = best_index_int8(16, kvalues_iq4nl_f, x1);
+
+        dst.qs[j] = xi0 | (xi1 << 4);
+
+        const float v0 = kvalues_iq4nl_f[xi0];
+        const float v1 = kvalues_iq4nl_f[xi1];
+        const float w0 = src[0        + j]*src[0        + j];
+        const float w1 = src[QK4_NL/2 + j]*src[QK4_NL/2 + j];
+        sumqx += w0*v0*src[j] + w1*v1*src[QK4_NL/2 + j];
+        sumq2 += w0*v0*v0 + w1*v1*v1;
+
+    }
+
+    dst.d = sumq2 > 0 ? sumqx/sumq2 : d;
+}

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

@@ -0,0 +1,389 @@
+#include "common.h"
+#include "dequantize.h"
+#include "quantize.h"
+
+template<typename T0, typename T1>
+kernel void kernel_cpy_t_t(
+        constant ggml_metal_kargs_cpy & args,
+        device  const char * src0,
+        device        char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    const int32_t i03 = tgpig[2];
+    const int32_t i02 = tgpig[1];
+    const int32_t i01 = ntg[1] == 1 ? tgpig[0]%args.ne01 : tgpig[0]*ntg[1] + tpitg.y;
+    const int32_t iw0 = ntg[1] == 1 ? tgpig[0]/args.ne01 : 0;
+
+    if (i01 >= args.ne01) {
+        return;
+    }
+
+    const int64_t n = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00;
+
+    const int32_t i3 = n/(args.ne2*args.ne1*args.ne0);
+    const int32_t i2 = (n - i3*args.ne2*args.ne1*args.ne0)/(args.ne1*args.ne0);
+    const int32_t i1 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0)/args.ne0;
+    const int32_t i0 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0 - i1*args.ne0);
+
+    device T1 * dst_data = (device T1 *) (dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
+
+    for (int32_t i00 = iw0*ntg[0] + tpitg.x; i00 < args.ne00;) {
+        device const T0 * src = (device T0 *)(src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + i00*args.nb00);
+        dst_data[i00] = (T1) src[0];
+        break;
+    }
+}
+
+typedef decltype(kernel_cpy_t_t<float, float>) kernel_cpy_t;
+
+template [[host_name("kernel_cpy_f32_f32")]]   kernel kernel_cpy_t kernel_cpy_t_t<float,   float>;
+template [[host_name("kernel_cpy_f32_f16")]]   kernel kernel_cpy_t kernel_cpy_t_t<float,   half>;
+template [[host_name("kernel_cpy_f32_i32")]]   kernel kernel_cpy_t kernel_cpy_t_t<float,   int32_t>;
+template [[host_name("kernel_cpy_i32_f32")]]   kernel kernel_cpy_t kernel_cpy_t_t<int32_t, float>;
+template [[host_name("kernel_cpy_i32_i32")]]   kernel kernel_cpy_t kernel_cpy_t_t<int32_t, int32_t>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_cpy_f32_bf16")]]  kernel kernel_cpy_t kernel_cpy_t_t<float,   bfloat>;
+#endif
+template [[host_name("kernel_cpy_f16_f32")]]   kernel kernel_cpy_t kernel_cpy_t_t<half,    float>;
+template [[host_name("kernel_cpy_f16_f16")]]   kernel kernel_cpy_t kernel_cpy_t_t<half,    half>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_cpy_bf16_f32")]]  kernel kernel_cpy_t kernel_cpy_t_t<bfloat,  float>;
+template [[host_name("kernel_cpy_bf16_bf16")]] kernel kernel_cpy_t kernel_cpy_t_t<bfloat,  bfloat>;
+#endif
+
+template<short QK,
+         typename block_q,
+         void (*quantize_func)(device const float *, device block_q &)>
+kernel void kernel_cpy_f32_q(
+        constant ggml_metal_kargs_cpy & args,
+        device const char * src0,
+        device char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    const int32_t i03 = tgpig[2];
+    const int32_t i02 = tgpig[1];
+    const int32_t i01 = ntg[1] == 1 ? tgpig[0]%args.ne01 : tgpig[0]*ntg[1] + tpitg.y;
+    const int32_t iw0 = ntg[1] == 1 ? tgpig[0]/args.ne01 : 0;
+
+    if (i01 >= args.ne01) {
+        return;
+    }
+
+    const int64_t n = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00;
+
+    const int32_t i3 = n / (args.ne2*args.ne1*args.ne0);
+    const int32_t i2 = (n - i3*args.ne2*args.ne1*args.ne0) / (args.ne1*args.ne0);
+    const int32_t i1 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0) / args.ne0;
+    const int32_t i0 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0 - i1*args.ne0)/QK;
+
+    device block_q * dst_data = (device block_q *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
+
+    for (int32_t i00 = iw0*ntg[0] + tpitg.x; i00 < args.nk0;) {
+        device const float * src = (device const float *)(src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + (i00*QK)*args.nb00);
+
+        quantize_func(src, dst_data[i00]);
+
+        break;
+    }
+}
+
+typedef decltype(kernel_cpy_f32_q<QK8_0,  block_q8_0,  quantize_q8_0>)  cpy_f_q_t;
+
+template [[host_name("kernel_cpy_f32_q8_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK8_0,  block_q8_0,   quantize_q8_0>;
+template [[host_name("kernel_cpy_f32_q1_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK1_0,  block_q1_0,   quantize_q1_0>;
+template [[host_name("kernel_cpy_f32_q4_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK4_0,  block_q4_0,   quantize_q4_0>;
+template [[host_name("kernel_cpy_f32_q4_1")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK4_1,  block_q4_1,   quantize_q4_1>;
+template [[host_name("kernel_cpy_f32_q5_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK5_0,  block_q5_0,   quantize_q5_0>;
+template [[host_name("kernel_cpy_f32_q5_1")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK5_1,  block_q5_1,   quantize_q5_1>;
+template [[host_name("kernel_cpy_f32_iq4_nl")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK4_NL, block_iq4_nl, quantize_iq4_nl>;
+
+template<typename T4x4, typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread T4x4 &)>
+kernel void kernel_cpy_q_f32(
+        constant ggml_metal_kargs_cpy & args,
+        device  const char * src0,
+        device        char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    const int32_t i03 = tgpig[2];
+    const int32_t i02 = tgpig[1];
+    const int32_t i01 = ntg[1] == 1 ? tgpig[0]%args.ne01 : tgpig[0]*ntg[1] + tpitg.y;
+    const int32_t iw0 = ntg[1] == 1 ? tgpig[0]/args.ne01 : 0;
+
+    if (i01 >= args.ne01) {
+        return;
+    }
+
+    const int64_t n = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00;
+
+    const int32_t i3 = n/(args.ne2*args.ne1*args.ne0);
+    const int32_t i2 = (n - i3*args.ne2*args.ne1*args.ne0)/(args.ne1*args.ne0);
+    const int32_t i1 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0)/args.ne0;
+    const int32_t i0 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0 - i1*args.ne0);
+
+    device const block_q * src_data = (device const block_q *)(src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
+    device       T4x4    * dst_data = (device       T4x4    *)(dst  +  i3*args.nb3  +  i2*args.nb2  +  i1*args.nb1 + i0*args.nb0);
+
+    for (int32_t i00 = iw0*ntg[0] + tpitg.x; i00 < args.nk0;) {
+        T4x4 temp;
+        dequantize_func(src_data + i00/nl, i00%nl, temp);
+        dst_data[i00] = temp;
+
+        break;
+    }
+}
+
+typedef decltype(kernel_cpy_q_f32<float4x4, block_q4_0, 2, dequantize_q4_0>) cpy_q_f_t;
+
+template [[host_name("kernel_cpy_q1_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q1_0, 8, dequantize_q1_0>;
+template [[host_name("kernel_cpy_q4_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q4_0, 2, dequantize_q4_0>;
+template [[host_name("kernel_cpy_q4_1_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q4_1, 2, dequantize_q4_1>;
+template [[host_name("kernel_cpy_q5_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q5_0, 2, dequantize_q5_0>;
+template [[host_name("kernel_cpy_q5_1_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q5_1, 2, dequantize_q5_1>;
+template [[host_name("kernel_cpy_q8_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q8_0, 2, dequantize_q8_0>;
+
+template [[host_name("kernel_cpy_q1_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q1_0, 8, dequantize_q1_0>;
+template [[host_name("kernel_cpy_q4_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q4_0, 2, dequantize_q4_0>;
+template [[host_name("kernel_cpy_q4_1_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q4_1, 2, dequantize_q4_1>;
+template [[host_name("kernel_cpy_q5_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q5_0, 2, dequantize_q5_0>;
+template [[host_name("kernel_cpy_q5_1_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q5_1, 2, dequantize_q5_1>;
+template [[host_name("kernel_cpy_q8_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q8_0, 2, dequantize_q8_0>;
+
+template<typename T>
+kernel void kernel_concat(
+        constant ggml_metal_kargs_concat & args,
+        device  const char * src0,
+        device  const char * src1,
+        device        char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+
+    const int i3 = tgpig.z;
+    const int i2 = tgpig.y;
+    const int i1 = ntg.y == 1 ? tgpig.x : tgpig.x*ntg.y + tpitg.y;
+
+    if (i1 >= args.ne1) {
+        return;
+    }
+
+    int o[4] = {0, 0, 0, 0};
+    o[args.dim] = args.dim == 0 ? args.ne00 : (args.dim == 1 ? args.ne01 : (args.dim == 2 ? args.ne02 : args.ne03));
+
+    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
+        device const T * x;
+
+        if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {
+            x = (device const T *)(src0 + (i3       )*args.nb03 + (i2       )*args.nb02 + (i1       )*args.nb01 + (i0       )*args.nb00);
+        } else {
+            x = (device const T *)(src1 + (i3 - o[3])*args.nb13 + (i2 - o[2])*args.nb12 + (i1 - o[1])*args.nb11 + (i0 - o[0])*args.nb10);
+        }
+
+        device T * y = (device T *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
+
+        *y = *x;
+    }
+}
+
+typedef decltype(kernel_concat<float>) kernel_concat_t;
+
+template [[host_name("kernel_concat_f32")]]  kernel kernel_concat_t kernel_concat<float>;
+template [[host_name("kernel_concat_f16")]]  kernel kernel_concat_t kernel_concat<half>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_concat_bf16")]] kernel kernel_concat_t kernel_concat<bfloat>;
+#endif
+template [[host_name("kernel_concat_i8")]]   kernel kernel_concat_t kernel_concat<char>;
+template [[host_name("kernel_concat_i16")]]  kernel kernel_concat_t kernel_concat<short>;
+template [[host_name("kernel_concat_i32")]]  kernel kernel_concat_t kernel_concat<int>;
+template [[host_name("kernel_concat_i64")]]  kernel kernel_concat_t kernel_concat<long>;
+
+template<typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread float4x4 &)>
+kernel void kernel_get_rows_q(
+        constant ggml_metal_kargs_get_rows & args,
+        device const void * src0,
+        device const void * src1,
+        device       void * dst,
+        uint3               tgpig[[threadgroup_position_in_grid]],
+        ushort              tiitg[[thread_index_in_threadgroup]],
+        ushort3             ntg  [[threads_per_threadgroup]]) {
+    const int32_t iw0 = tgpig.x/args.ne10;
+    const int32_t i10 = tgpig.x%args.ne10;
+    const int32_t i11 = tgpig.y;
+    const int32_t i12 = tgpig.z;
+
+    const int32_t r = ((const device int32_t *) ((const device char *) src1 + i12*args.nb12 + i11*args.nb11 + i10*args.nb10))[0];
+
+    const int32_t i02 = i11;
+    const int32_t i03 = i12;
+
+    auto psrc = (device const block_q *) ((const device char *) src0 + i03*args.nb03 + i02*args.nb02 +   r*args.nb01);
+    auto pdst = (device      float4x4 *) ((      device char *) dst  + i12*args.nb3  + i11*args.nb2  + i10*args.nb1);
+
+    for (int ind = iw0*ntg.x + tiitg; ind < args.ne00t;) {
+        float4x4 temp;
+        dequantize_func(psrc + ind/nl, ind%nl, temp);
+        pdst[ind] = temp;
+
+        break;
+    }
+}
+
+template<typename T0, typename T>
+kernel void kernel_get_rows_f(
+        constant ggml_metal_kargs_get_rows & args,
+        device const void * src0,
+        device const void * src1,
+        device       void * dst,
+        uint3               tgpig[[threadgroup_position_in_grid]],
+        ushort              tiitg[[thread_index_in_threadgroup]],
+        ushort3             ntg [[threads_per_threadgroup]]) {
+    const int32_t iw0 = tgpig.x/args.ne10;
+    const int32_t i10 = tgpig.x%args.ne10;
+    const int32_t i11 = tgpig.y;
+    const int32_t i12 = tgpig.z;
+
+    const int32_t r = ((const device int32_t *) ((const device char *) src1 + i12*args.nb12 + i11*args.nb11 + i10*args.nb10))[0];
+
+    const int32_t i02 = i11;
+    const int32_t i03 = i12;
+
+    auto psrc = (const device T0 *) ((const device char *) src0 + i03*args.nb03 + i02*args.nb02 +   r*args.nb01);
+    auto pdst = (      device T  *) ((      device char *)  dst + i12*args.nb3  + i11*args.nb2  + i10*args.nb1);
+
+    for (int ind = iw0*ntg.x + tiitg; ind < args.ne00t;) {
+        pdst[ind] = psrc[ind];
+
+        break;
+    }
+}
+
+typedef decltype(kernel_get_rows_f<float, float>) get_rows_f_t;
+
+template [[host_name("kernel_get_rows_f32")]]  kernel get_rows_f_t kernel_get_rows_f<float, float>;
+template [[host_name("kernel_get_rows_f16")]]  kernel get_rows_f_t kernel_get_rows_f<half,  float>;
+template [[host_name("kernel_get_rows_i32")]]  kernel get_rows_f_t kernel_get_rows_f<int32_t, int32_t>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_get_rows_bf16")]] kernel get_rows_f_t kernel_get_rows_f<bfloat, float>;
+#endif
+
+typedef decltype(kernel_get_rows_q<block_q4_0, 2, dequantize_q4_0>) get_rows_q_t;
+
+template [[host_name("kernel_get_rows_q1_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q1_0,    8, dequantize_q1_0>;
+template [[host_name("kernel_get_rows_q4_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_0,    2, dequantize_q4_0>;
+template [[host_name("kernel_get_rows_q4_1")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_1,    2, dequantize_q4_1>;
+template [[host_name("kernel_get_rows_q5_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_0,    2, dequantize_q5_0>;
+template [[host_name("kernel_get_rows_q5_1")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_1,    2, dequantize_q5_1>;
+template [[host_name("kernel_get_rows_q8_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q8_0,    2, dequantize_q8_0>;
+template [[host_name("kernel_get_rows_mxfp4")]]   kernel get_rows_q_t kernel_get_rows_q<block_mxfp4,   2, dequantize_mxfp4>;
+template [[host_name("kernel_get_rows_q2_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q2_K,    QK_NL, dequantize_q2_K>;
+template [[host_name("kernel_get_rows_q3_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q3_K,    QK_NL, dequantize_q3_K>;
+template [[host_name("kernel_get_rows_q4_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_K,    QK_NL, dequantize_q4_K>;
+template [[host_name("kernel_get_rows_q5_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_K,    QK_NL, dequantize_q5_K>;
+template [[host_name("kernel_get_rows_q6_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q6_K,    QK_NL, dequantize_q6_K>;
+template [[host_name("kernel_get_rows_iq2_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq2_xxs, QK_NL, dequantize_iq2_xxs>;
+template [[host_name("kernel_get_rows_iq2_xs")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq2_xs,  QK_NL, dequantize_iq2_xs>;
+template [[host_name("kernel_get_rows_iq3_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq3_xxs, QK_NL, dequantize_iq3_xxs>;
+template [[host_name("kernel_get_rows_iq3_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq3_s,   QK_NL, dequantize_iq3_s>;
+template [[host_name("kernel_get_rows_iq2_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq2_s,   QK_NL, dequantize_iq2_s>;
+template [[host_name("kernel_get_rows_iq1_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq1_s,   QK_NL, dequantize_iq1_s>;
+template [[host_name("kernel_get_rows_iq1_m")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq1_m,   QK_NL, dequantize_iq1_m>;
+template [[host_name("kernel_get_rows_iq4_nl")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq4_nl,  2,     dequantize_iq4_nl>;
+template [[host_name("kernel_get_rows_iq4_xs")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq4_xs,  QK_NL, dequantize_iq4_xs>;
+
+template<typename TS, typename TI, typename block_q, void (*quantize_func)(device const float *, device block_q &)>
+kernel void kernel_set_rows_q32(
+        constant ggml_metal_kargs_set_rows & args,
+        device const  void * src0,
+        device const  void * src1,
+        device       float * dst,
+        uint3                tgpig[[threadgroup_position_in_grid]],
+        uint                 tiitg[[thread_index_in_threadgroup]],
+        uint3                tptg [[threads_per_threadgroup]]) {
+    const int32_t i03 = tgpig.z;
+    const int32_t i02 = tgpig.y;
+
+    const int32_t i12 = i03%args.ne12;
+    const int32_t i11 = i02%args.ne11;
+
+    const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x;
+    if (i01 >= args.ne01) {
+        return;
+    }
+
+    const int32_t i10 = i01;
+    const TI      i1  = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0];
+
+          device block_q * dst_row = (      device block_q *) ((      device char *) dst  +  i1*args.nb1  + i02*args.nb2  + i03*args.nb3);
+    const device TS      * src_row = (const device TS      *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
+
+    for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) {
+        quantize_func(src_row + 32*ind, dst_row[ind]);
+    }
+}
+
+template<typename TS, typename TI, typename TD>
+kernel void kernel_set_rows_f(
+        constant ggml_metal_kargs_set_rows & args,
+        device const  void * src0,
+        device const  void * src1,
+        device       float * dst,
+        uint3                tgpig[[threadgroup_position_in_grid]],
+        uint                 tiitg[[thread_index_in_threadgroup]],
+        uint3                tptg [[threads_per_threadgroup]]) {
+    const int32_t i03 = tgpig.z;
+    const int32_t i02 = tgpig.y;
+
+    const int32_t i12 = i03%args.ne12;
+    const int32_t i11 = i02%args.ne11;
+
+    const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x;
+    if (i01 >= args.ne01) {
+        return;
+    }
+
+    const int32_t i10 = i01;
+    const TI      i1  = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0];
+
+          device TD * dst_row = (      device TD *) ((      device char *) dst  +  i1*args.nb1  + i02*args.nb2  + i03*args.nb3);
+    const device TS * src_row = (const device TS *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
+
+    for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) {
+        dst_row[ind] = (TD) src_row[ind];
+    }
+}
+
+typedef decltype(kernel_set_rows_f<float, int64_t, float>) set_rows_f_t;
+
+template [[host_name("kernel_set_rows_f32_i64_f32")]]   kernel set_rows_f_t kernel_set_rows_f<float, int64_t, float>;
+template [[host_name("kernel_set_rows_f32_i32_f32")]]   kernel set_rows_f_t kernel_set_rows_f<float, int32_t, float>;
+template [[host_name("kernel_set_rows_f32_i64_f16")]]   kernel set_rows_f_t kernel_set_rows_f<float, int64_t, half>;
+template [[host_name("kernel_set_rows_f32_i32_f16")]]   kernel set_rows_f_t kernel_set_rows_f<float, int32_t, half>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_set_rows_f32_i64_bf16")]]  kernel set_rows_f_t kernel_set_rows_f<float, int64_t, bfloat>;
+template [[host_name("kernel_set_rows_f32_i32_bf16")]]  kernel set_rows_f_t kernel_set_rows_f<float, int32_t, bfloat>;
+#endif
+
+template [[host_name("kernel_set_rows_f16_i64_f16")]]   kernel set_rows_f_t kernel_set_rows_f<half, int64_t, half>;
+template [[host_name("kernel_set_rows_f16_i32_f16")]]   kernel set_rows_f_t kernel_set_rows_f<half, int32_t, half>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_set_rows_bf16_i64_bf16")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int64_t, bfloat>;
+template [[host_name("kernel_set_rows_bf16_i32_bf16")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int32_t, bfloat>;
+#endif
+
+typedef decltype(kernel_set_rows_q32<float, int64_t, block_q8_0, quantize_q8_0>) set_rows_q32_t;
+
+template [[host_name("kernel_set_rows_f32_i64_q8_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q8_0,   quantize_q8_0>;
+template [[host_name("kernel_set_rows_f32_i32_q8_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q8_0,   quantize_q8_0>;
+template [[host_name("kernel_set_rows_f32_i64_q4_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q4_0,   quantize_q4_0>;
+template [[host_name("kernel_set_rows_f32_i32_q4_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q4_0,   quantize_q4_0>;
+template [[host_name("kernel_set_rows_f32_i64_q4_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q4_1,   quantize_q4_1>;
+template [[host_name("kernel_set_rows_f32_i32_q4_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q4_1,   quantize_q4_1>;
+template [[host_name("kernel_set_rows_f32_i64_q5_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q5_0,   quantize_q5_0>;
+template [[host_name("kernel_set_rows_f32_i32_q5_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q5_0,   quantize_q5_0>;
+template [[host_name("kernel_set_rows_f32_i64_q5_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q5_1,   quantize_q5_1>;
+template [[host_name("kernel_set_rows_f32_i32_q5_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q5_1,   quantize_q5_1>;
+template [[host_name("kernel_set_rows_f32_i64_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_iq4_nl, quantize_iq4_nl>;
+template [[host_name("kernel_set_rows_f32_i32_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_iq4_nl, quantize_iq4_nl>;
+

ggml/src/ggml-metal/kernels/reduce.metal

@@ -0,0 +1,228 @@
+#include "common.h"
+
+kernel void kernel_op_sum_f32(
+        constant ggml_metal_kargs_sum & args,
+        device const float * src0,
+        device       float * dst,
+        threadgroup  float * shmem_f32 [[threadgroup(0)]],
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort  tiisg[[thread_index_in_simdgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+
+    if (args.np == 0) {
+        return;
+    }
+
+    // TODO: become function constant
+    const uint nsg = (ntg.x + 31) / 32;
+
+    float sumf = 0;
+
+    for (uint64_t i0 = tpitg.x; i0 < args.np; i0 += ntg.x) {
+        sumf += src0[i0];
+    }
+
+    sumf = simd_sum(sumf);
+
+    if (tiisg == 0) {
+        shmem_f32[sgitg] = sumf;
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    float total = 0;
+
+    if (sgitg == 0) {
+        float v = 0;
+
+        if (tpitg.x < nsg) {
+            v = shmem_f32[tpitg.x];
+        }
+
+        total = simd_sum(v);
+
+        if (tpitg.x == 0) {
+            dst[0] = total;
+        }
+    }
+}
+
+constant short FC_sum_rows_op [[function_constant(FC_SUM_ROWS + 0)]];
+
+template <typename T0, typename T>
+kernel void kernel_sum_rows_impl(
+        constant ggml_metal_kargs_sum_rows & args,
+        device const char * src0,
+        device       char * dst,
+        threadgroup  char * shmem [[threadgroup(0)]],
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort  tiisg[[thread_index_in_simdgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+#define FC_OP  FC_sum_rows_op
+
+    const int i3 = tgpig.z;
+    const int i2 = tgpig.y;
+    const int i1 = tgpig.x;
+
+    threadgroup T0 * shmem_t = (threadgroup T0 *) shmem;
+
+    if (sgitg == 0) {
+        shmem_t[tiisg] = 0.0f;
+    }
+
+    device const T0 * src_row = (device const T0 *) (src0 + i1*args.nb01 + i2*args.nb02 + i3*args.nb03);
+    device       T  * dst_row = (device       T  *) (dst  + i1*args.nb1  + i2*args.nb2  + i3*args.nb3);
+
+    T0 sumf = T0(0.0f);
+
+    for (int64_t i0 = tpitg.x; i0 < args.ne00; i0 += ntg.x) {
+        sumf += src_row[i0];
+    }
+
+    sumf = simd_sum(sumf);
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    if (tiisg == 0) {
+        shmem_t[sgitg] = sumf;
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    sumf = shmem_t[tiisg];
+    sumf = simd_sum(sumf);
+
+    if (tpitg.x == 0) {
+        if (FC_OP == OP_SUM_ROWS_NUM_MEAN) {
+            if (is_same<float4, T0>::value) {
+                dst_row[0] = sum(sumf) / (4*args.ne00);
+            } else {
+                dst_row[0] = sum(sumf) / args.ne00;
+            }
+        } else {
+            dst_row[0] = sum(sumf);
+        }
+    }
+
+#undef FC_OP
+}
+
+typedef decltype(kernel_sum_rows_impl<float, float>) kernel_sum_rows_t;
+
+template [[host_name("kernel_sum_rows_f32_f32")]]   kernel kernel_sum_rows_t kernel_sum_rows_impl<float,  float>;
+template [[host_name("kernel_sum_rows_f32_f32_4")]] kernel kernel_sum_rows_t kernel_sum_rows_impl<float4, float>;
+
+template<typename T>
+kernel void kernel_cumsum_blk(
+        constant ggml_metal_kargs_cumsum_blk & args,
+        device const char * src0,
+        device       char * tmp,
+        device       char * dst,
+        threadgroup  char * shmem [[threadgroup(0)]],
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort  tiisg[[thread_index_in_simdgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    const int ib = tgpig[0]/args.ne01;
+
+    const int i00 = ib*ntg.x;
+    const int i01 = tgpig[0]%args.ne01;
+    const int i02 = tgpig[1];
+    const int i03 = tgpig[2];
+
+    device const float * src0_row = (device const float *) (src0 +
+            args.nb01*i01 +
+            args.nb02*i02 +
+            args.nb03*i03);
+
+    threadgroup float * shmem_f32 = (threadgroup float *) shmem;
+
+    float v = 0.0f;
+
+    if (i00 + tpitg.x < args.ne00) {
+        v = src0_row[i00 + tpitg.x];
+    }
+
+    float s = simd_prefix_inclusive_sum(v);
+
+    if (tiisg == N_SIMDWIDTH - 1) {
+        shmem_f32[sgitg] = s;
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    if (sgitg == 0) {
+        shmem_f32[tiisg] = simd_prefix_exclusive_sum(shmem_f32[tiisg]);
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    s += shmem_f32[sgitg];
+
+    device float * dst_row = (device float *) dst +
+        args.ne00*i01 +
+        args.ne00*args.ne01*i02 +
+        args.ne00*args.ne01*args.ne02*i03;
+
+    if (i00 + tpitg.x < args.ne00) {
+        dst_row[i00 + tpitg.x] = s;
+    }
+
+    if (args.outb && tpitg.x == ntg.x - 1) {
+        device float * tmp_row = (device float *) tmp +
+            args.net0*i01 +
+            args.net0*args.net1*i02 +
+            args.net0*args.net1*args.net2*i03;
+
+        tmp_row[ib] = s;
+    }
+}
+
+typedef decltype(kernel_cumsum_blk<float>) kernel_cumsum_blk_t;
+
+template [[host_name("kernel_cumsum_blk_f32")]] kernel kernel_cumsum_blk_t kernel_cumsum_blk<float>;
+
+template<typename T>
+kernel void kernel_cumsum_add(
+        constant ggml_metal_kargs_cumsum_add & args,
+        device const char * tmp,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort  tiisg[[thread_index_in_simdgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    const int ib = tgpig[0]/args.ne01;
+
+    if (ib == 0) {
+        return;
+    }
+
+    const int i00 = ib*ntg.x;
+    const int i01 = tgpig[0]%args.ne01;
+    const int i02 = tgpig[1];
+    const int i03 = tgpig[2];
+
+    device const float * tmp_row = (device const float *) (tmp +
+            args.nbt1*i01 +
+            args.nbt2*i02 +
+            args.nbt3*i03);
+
+    device float * dst_row = (device float *) dst +
+        args.ne00*i01 +
+        args.ne00*args.ne01*i02 +
+        args.ne00*args.ne01*args.ne02*i03;
+
+    if (i00 + tpitg.x < args.ne00) {
+        dst_row[i00 + tpitg.x] += tmp_row[ib - 1];
+    }
+}
+
+typedef decltype(kernel_cumsum_add<float>) kernel_cumsum_add_t;
+
+template [[host_name("kernel_cumsum_add_f32")]] kernel kernel_cumsum_add_t kernel_cumsum_add<float>;

ggml/src/ggml-metal/kernels/rope.metal

@@ -0,0 +1,318 @@
+#include "common.h"
+
+constant bool FC_rope_is_imrope [[function_constant(FC_ROPE + 0)]];
+constant bool FC_rope_is_back   [[function_constant(FC_ROPE + 1)]];
+
+static float rope_yarn_ramp(const float low, const float high, const int i0) {
+    const float y = (i0 / 2 - low) / max(0.001f, high - low);
+    return 1.0f - min(1.0f, max(0.0f, y));
+}
+
+// YaRN algorithm based on LlamaYaRNScaledRotaryEmbedding.py from https://github.com/jquesnelle/yarn
+// MIT licensed. Copyright (c) 2023 Jeffrey Quesnelle and Bowen Peng.
+static void rope_yarn(
+    float theta_extrap, float freq_scale, float corr_dims[2], int i0, float ext_factor, float mscale,
+    thread float * cos_theta, thread float * sin_theta) {
+    // Get n-d rotational scaling corrected for extrapolation
+    float theta_interp = freq_scale * theta_extrap;
+    float theta = theta_interp;
+    if (ext_factor != 0.0f) {
+        float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], i0) * ext_factor;
+        theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix;
+
+        // Get n-d magnitude scaling corrected for interpolation
+        mscale *= 1.0f + 0.1f * log(1.0f / freq_scale);
+    }
+    *cos_theta = cos(theta) * mscale;
+    *sin_theta = sin(theta) * mscale;
+    if (FC_rope_is_back) {
+        *sin_theta *= -1.0f;
+    }
+}
+
+// Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get
+// `corr_fac(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))`
+static float rope_yarn_corr_factor(int n_dims, int n_ctx_orig, float n_rot, float base) {
+    return n_dims * log(n_ctx_orig / (n_rot * 2 * M_PI_F)) / (2 * log(base));
+}
+
+static void rope_yarn_corr_dims(
+    int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2]
+) {
+    // start and end correction dims
+    dims[0] = max(0.0f,         floor(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_fast, freq_base)));
+    dims[1] = min(n_dims - 1.0f, ceil(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_slow, freq_base)));
+}
+
+template<typename T>
+kernel void kernel_rope_norm(
+        constant ggml_metal_kargs_rope & args,
+        device const char * src0,
+        device const char * src1,
+        device const char * src2,
+        device       char * dst,
+        ushort  tiitg[[thread_index_in_threadgroup]],
+        ushort3 tptg [[threads_per_threadgroup]],
+        uint3   tgpig[[threadgroup_position_in_grid]]) {
+    const int i3 = tgpig[2];
+    const int i2 = tgpig[1];
+    const int i1 = tgpig[0];
+
+    float corr_dims[2];
+    rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
+
+    device const int32_t * pos = (device const int32_t *) src1;
+
+    const float theta_base = (float) pos[i2];
+    const float inv_ndims = -1.f/args.n_dims;
+
+    float cos_theta;
+    float sin_theta;
+
+    for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
+        if (i0 < args.n_dims) {
+            const int ic = i0/2;
+
+            const float theta = theta_base * pow(args.freq_base, inv_ndims*i0);
+
+            const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
+
+            rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
+
+            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
+            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + i0*args.nb0);
+
+            const float x0 = src[0];
+            const float x1 = src[1];
+
+            dst_data[0] = x0*cos_theta - x1*sin_theta;
+            dst_data[1] = x0*sin_theta + x1*cos_theta;
+        } else {
+            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
+            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + i0*args.nb0);
+
+            dst_data[0] = src[0];
+            dst_data[1] = src[1];
+        }
+    }
+}
+
+template<typename T>
+kernel void kernel_rope_neox(
+        constant ggml_metal_kargs_rope & args,
+        device const char * src0,
+        device const char * src1,
+        device const char * src2,
+        device       char * dst,
+        ushort  tiitg[[thread_index_in_threadgroup]],
+        ushort3 tptg [[threads_per_threadgroup]],
+        uint3   tgpig[[threadgroup_position_in_grid]]) {
+    const int i3 = tgpig[2];
+    const int i2 = tgpig[1];
+    const int i1 = tgpig[0];
+
+    float corr_dims[2];
+    rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
+
+    device const int32_t * pos = (device const int32_t *) src1;
+
+    const float theta_base = (float) pos[i2];
+    const float inv_ndims = -1.f/args.n_dims;
+
+    float cos_theta;
+    float sin_theta;
+
+    for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
+        if (i0 < args.n_dims) {
+            const int ic = i0/2;
+
+            const float theta = theta_base * pow(args.freq_base, inv_ndims*i0);
+
+            const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
+
+            rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
+
+            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + ic*args.nb00);
+            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + ic*args.nb0);
+
+            const float x0 = src[0];
+            const float x1 = src[args.n_dims/2];
+
+            dst_data[0]             = x0*cos_theta - x1*sin_theta;
+            dst_data[args.n_dims/2] = x0*sin_theta + x1*cos_theta;
+        } else {
+            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
+            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + i0*args.nb0);
+
+            dst_data[0] = src[0];
+            dst_data[1] = src[1];
+        }
+    }
+}
+
+template<typename T>
+kernel void kernel_rope_multi(
+        constant ggml_metal_kargs_rope & args,
+        device const char * src0,
+        device const char * src1,
+        device const char * src2,
+        device       char * dst,
+        ushort  tiitg[[thread_index_in_threadgroup]],
+        ushort3 tptg [[threads_per_threadgroup]],
+        uint3   tgpig[[threadgroup_position_in_grid]]) {
+    const int i3 = tgpig[2];
+    const int i2 = tgpig[1];
+    const int i1 = tgpig[0];
+
+    float corr_dims[2];
+    rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
+
+    device const int32_t * pos = (device const int32_t *) src1;
+
+    const float inv_ndims = -1.f/args.n_dims;
+
+    float cos_theta;
+    float sin_theta;
+
+    for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
+        if (i0 < args.n_dims) {
+            const int ic = i0/2;
+
+            // mrope theta calculations
+            // note: the rest is the same as kernel_rope_neox
+            const int sect_dims = args.sect_0 + args.sect_1 + args.sect_2 + args.sect_3;
+            const int sec_w01   = args.sect_0 + args.sect_1;               // end of section 1
+            const int sec_w012  = args.sect_0 + args.sect_1 + args.sect_2; // end of section 2
+            const int sector    = ic % sect_dims;
+
+            float theta_base;
+            if (FC_rope_is_imrope) {
+                if (sector % 3 == 1 && sector < 3 * args.sect_1) { // h
+                    theta_base = (float) pos[i2 + args.ne02 * 1];
+                } else if (sector % 3 == 2 && sector < 3 * args.sect_2) { // w
+                    theta_base = (float) pos[i2 + args.ne02 * 2];
+                } else if (sector % 3 == 0 && sector < 3 * args.sect_0) { // t
+                    theta_base = (float) pos[i2 + args.ne02 * 0];
+                } else { // e
+                    theta_base = (float) pos[i2 + args.ne02 * 3];
+                }
+            } else {
+                if (sector < args.sect_0) {
+                    theta_base = (float) pos[i2];
+                } else if (sector < sec_w01) {
+                    theta_base = (float) pos[i2 + args.ne02 * 1];
+                } else if (sector < sec_w012) {
+                    theta_base = (float) pos[i2 + args.ne02 * 2];
+                } else {
+                    theta_base = (float) pos[i2 + args.ne02 * 3];
+                }
+            }
+            // end of mrope
+
+            const float theta = theta_base * pow(args.freq_base, inv_ndims*i0);
+
+            const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
+
+            rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
+
+            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + ic*args.nb00);
+            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + ic*args.nb0);
+
+            const float x0 = src[0];
+            const float x1 = src[args.n_dims/2];
+
+            dst_data[0]             = x0*cos_theta - x1*sin_theta;
+            dst_data[args.n_dims/2] = x0*sin_theta + x1*cos_theta;
+        } else {
+            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
+            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + i0*args.nb0);
+
+            dst_data[0] = src[0];
+            dst_data[1] = src[1];
+        }
+    }
+}
+
+template<typename T>
+kernel void kernel_rope_vision(
+        constant ggml_metal_kargs_rope & args,
+        device const char * src0,
+        device const char * src1,
+        device const char * src2,
+        device       char * dst,
+        ushort  tiitg[[thread_index_in_threadgroup]],
+        ushort3 tptg [[threads_per_threadgroup]],
+        uint3   tgpig[[threadgroup_position_in_grid]]) {
+    const int i3 = tgpig[2];
+    const int i2 = tgpig[1];
+    const int i1 = tgpig[0];
+
+    float corr_dims[2];
+    rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
+
+    device const int32_t * pos = (device const int32_t *) src1;
+
+    const float inv_ndims = -1.f/args.n_dims;
+
+    float cos_theta;
+    float sin_theta;
+
+    for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
+        if (i0 < 2*args.n_dims) { // different from kernel_rope_multi
+            const int ic = i0/2;
+
+            // mrope theta calculations (only support 2 dimensions)
+            const int sect_dims = args.sect_0 + args.sect_1;
+            const int sector    = ic % sect_dims;
+
+            float p;
+            float theta_base;
+            if (sector < args.sect_1) {
+                p = (float) sector;
+                theta_base = (float) pos[i2];
+            } else {
+                p = (float) sector - args.sect_0;
+                theta_base = (float) pos[i2 + args.ne02];
+            }
+
+            const float theta = theta_base * pow(args.freq_base, 2.0f * inv_ndims * p);
+            // end of mrope
+
+            const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
+
+            rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
+
+            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + ic*args.nb00);
+            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + ic*args.nb0);
+
+            const float x0 = src[0];
+            const float x1 = src[args.n_dims]; // different from kernel_rope_multi
+
+            dst_data[0]           = x0*cos_theta - x1*sin_theta;
+            dst_data[args.n_dims] = x0*sin_theta + x1*cos_theta; // different from kernel_rope_multi
+        } else {
+            device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
+            device       T * dst_data  = (device T *)( dst + i3*args.nb3  + i2*args.nb2  + i1*args.nb1  + i0*args.nb0);
+
+            dst_data[0] = src[0];
+            dst_data[1] = src[1];
+        }
+    }
+}
+
+typedef decltype(kernel_rope_norm<float>) kernel_rope_norm_t;
+typedef decltype(kernel_rope_neox<float>) kernel_rope_neox_t;
+typedef decltype(kernel_rope_multi<float>) kernel_rope_multi_t;
+typedef decltype(kernel_rope_vision<float>) kernel_rope_vision_t;
+
+template [[host_name("kernel_rope_norm_f32")]] kernel kernel_rope_norm_t kernel_rope_norm<float>;
+template [[host_name("kernel_rope_norm_f16")]] kernel kernel_rope_norm_t kernel_rope_norm<half>;
+
+template [[host_name("kernel_rope_neox_f32")]] kernel kernel_rope_neox_t kernel_rope_neox<float>;
+template [[host_name("kernel_rope_neox_f16")]] kernel kernel_rope_neox_t kernel_rope_neox<half>;
+
+template [[host_name("kernel_rope_multi_f32")]] kernel kernel_rope_multi_t kernel_rope_multi<float>;
+template [[host_name("kernel_rope_multi_f16")]] kernel kernel_rope_multi_t kernel_rope_multi<half>;
+
+template [[host_name("kernel_rope_vision_f32")]] kernel kernel_rope_vision_t kernel_rope_vision<float>;
+template [[host_name("kernel_rope_vision_f16")]] kernel kernel_rope_vision_t kernel_rope_vision<half>;

ggml/src/ggml-metal/kernels/softmax.metal

@@ -0,0 +1,223 @@
+#include "common.h"
+
+template<typename T>
+kernel void kernel_soft_max(
+        constant ggml_metal_kargs_soft_max & args,
+        device const  char * src0,
+        device const  char * src1,
+        device const  char * src2,
+        device        char * dst,
+        threadgroup  float * buf [[threadgroup(0)]],
+        uint3 tgpig[[threadgroup_position_in_grid]],
+        uint3 tpitg[[thread_position_in_threadgroup]],
+        uint  sgitg[[simdgroup_index_in_threadgroup]],
+        uint  tiisg[[thread_index_in_simdgroup]],
+        uint3  tptg[[threads_per_threadgroup]]) {
+    const int32_t i03 = tgpig.z;
+    const int32_t i02 = tgpig.y;
+    const int32_t i01 = tgpig.x;
+
+    const int32_t i13 = i03%args.ne13;
+    const int32_t i12 = i02%args.ne12;
+    const int32_t i11 = i01;
+
+    device const float * psrc0 =                (device const float *) (src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
+    device const     T * pmask = src1 != src0 ? (device const T *    ) (src1 + i11*args.nb11 + i12*args.nb12 + i13*args.nb13) : nullptr;
+    device const float * psrc2 = src2 != src0 ? (device const float *) (src2)                                                 : nullptr;
+    device       float * pdst  =                (device       float *) (dst  + i01*args.nb1  + i02*args.nb2  + i03*args.nb3);
+
+    float slope = 1.0f;
+
+    // ALiBi
+    if (args.max_bias > 0.0f) {
+        const int32_t h = i02;
+
+        const float base = h < args.n_head_log2 ? args.m0 : args.m1;
+        const int   exp  = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
+
+        slope = pow(base, exp);
+    }
+
+    // parallel max
+    float lmax = psrc2 ? psrc2[i02] : -INFINITY;
+
+    for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
+        lmax = MAX(lmax, psrc0[i00]*args.scale + (pmask ? slope*pmask[i00] : 0.0f));
+    }
+
+    // find the max value in the block
+    float max_val = simd_max(lmax);
+    if (tptg.x > N_SIMDWIDTH) {
+        if (sgitg == 0) {
+            buf[tiisg] = -INFINITY;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiisg == 0) {
+            buf[sgitg] = max_val;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        max_val = buf[tiisg];
+        max_val = simd_max(max_val);
+    }
+
+    // parallel sum
+    float lsum = 0.0f;
+    for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
+        const float exp_psrc0 = exp((psrc0[i00]*args.scale + (pmask ? slope*pmask[i00] : 0.0f)) - max_val);
+        lsum += exp_psrc0;
+        pdst[i00] = exp_psrc0;
+    }
+
+    // This barrier fixes a failing test
+    // ref: https://github.com/ggml-org/ggml/pull/621#discussion_r1425156335
+    threadgroup_barrier(mem_flags::mem_none);
+
+    float sum = simd_sum(lsum);
+
+    if (tptg.x > N_SIMDWIDTH) {
+        if (sgitg == 0) {
+            buf[tiisg] = 0.0f;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiisg == 0) {
+            buf[sgitg] = sum;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        sum = buf[tiisg];
+        sum = simd_sum(sum);
+    }
+
+    if (psrc2) {
+        sum += exp(psrc2[i02] - max_val);
+    }
+
+    const float inv_sum = 1.0f/sum;
+
+    for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
+        pdst[i00] *= inv_sum;
+    }
+}
+
+template<typename T>
+kernel void kernel_soft_max_4(
+        constant ggml_metal_kargs_soft_max & args,
+        device const  char * src0,
+        device const  char * src1,
+        device const  char * src2,
+        device        char * dst,
+        threadgroup  float * buf [[threadgroup(0)]],
+        uint3 tgpig[[threadgroup_position_in_grid]],
+        uint3 tpitg[[thread_position_in_threadgroup]],
+        uint  sgitg[[simdgroup_index_in_threadgroup]],
+        uint  tiisg[[thread_index_in_simdgroup]],
+        uint3  tptg[[threads_per_threadgroup]]) {
+    const int32_t i03 = tgpig.z;
+    const int32_t i02 = tgpig.y;
+    const int32_t i01 = tgpig.x;
+
+    const int32_t i13 = i03%args.ne13;
+    const int32_t i12 = i02%args.ne12;
+    const int32_t i11 = i01;
+
+    device const float4 * psrc4 =                (device const float4 *) (src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
+    device const      T * pmask = src1 != src0 ? (device const T *     ) (src1 + i11*args.nb11 + i12*args.nb12 + i13*args.nb13) : nullptr;
+    device const float *  psrc2 = src2 != src0 ? (device const float * ) (src2)                                                 : nullptr;
+    device       float4 * pdst4 =                (device       float4 *) (dst  + i01*args.nb1  + i02*args.nb2  + i03*args.nb3);
+
+    float slope = 1.0f;
+
+    if (args.max_bias > 0.0f) {
+        const int32_t h = i02;
+
+        const float base = h < args.n_head_log2 ? args.m0 : args.m1;
+        const int   exp  = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
+
+        slope = pow(base, exp);
+    }
+
+    // parallel max
+    float4 lmax4 = psrc2 ? psrc2[i02] : -INFINITY;
+
+    for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
+        lmax4 = fmax(lmax4, psrc4[i00]*args.scale + (float4)((pmask ? slope*pmask[i00] : 0.0f)));
+    }
+
+    const float lmax = MAX(MAX(lmax4[0], lmax4[1]), MAX(lmax4[2], lmax4[3]));
+
+    float max_val = simd_max(lmax);
+    if (tptg.x > N_SIMDWIDTH) {
+        if (sgitg == 0) {
+            buf[tiisg] = -INFINITY;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiisg == 0) {
+            buf[sgitg] = max_val;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        max_val = buf[tiisg];
+        max_val = simd_max(max_val);
+    }
+
+    // parallel sum
+    float4 lsum4 = 0.0f;
+    for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
+        const float4 exp_psrc4 = exp((psrc4[i00]*args.scale + (float4)((pmask ? slope*pmask[i00] : 0.0f))) - max_val);
+        lsum4 += exp_psrc4;
+        pdst4[i00] = exp_psrc4;
+    }
+
+    const float lsum = lsum4[0] + lsum4[1] + lsum4[2] + lsum4[3];
+
+    // This barrier fixes a failing test
+    // ref: https://github.com/ggml-org/ggml/pull/621#discussion_r1425156335
+    threadgroup_barrier(mem_flags::mem_none);
+
+    float sum = simd_sum(lsum);
+
+    if (tptg.x > N_SIMDWIDTH) {
+        if (sgitg == 0) {
+            buf[tiisg] = 0.0f;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiisg == 0) {
+            buf[sgitg] = sum;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        sum = buf[tiisg];
+        sum = simd_sum(sum);
+    }
+
+    if (psrc2) {
+        sum += exp(psrc2[i02] - max_val);
+    }
+
+    const float inv_sum = 1.0f/sum;
+
+    for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
+        pdst4[i00] *= inv_sum;
+    }
+}
+
+typedef decltype(kernel_soft_max<float>)    kernel_soft_max_t;
+typedef decltype(kernel_soft_max_4<float4>) kernel_soft_max_4_t;
+
+template [[host_name("kernel_soft_max_f16")]]   kernel kernel_soft_max_t   kernel_soft_max<half>;
+template [[host_name("kernel_soft_max_f32")]]   kernel kernel_soft_max_t   kernel_soft_max<float>;
+template [[host_name("kernel_soft_max_f16_4")]] kernel kernel_soft_max_4_t kernel_soft_max_4<half4>;
+template [[host_name("kernel_soft_max_f32_4")]] kernel kernel_soft_max_4_t kernel_soft_max_4<float4>;

ggml/src/ggml-metal/kernels/solve_tri.metal

@@ -0,0 +1,75 @@
+#include "common.h"
+
+constant short FC_solve_tri_nsg [[function_constant(FC_SOLVE_TRI + 0)]];
+constant short FC_solve_tri_n   [[function_constant(FC_SOLVE_TRI + 1)]];
+constant short FC_solve_tri_k   [[function_constant(FC_SOLVE_TRI + 2)]];
+
+kernel void kernel_solve_tri_f32(
+        constant ggml_metal_kargs_solve_tri & args,
+        device   const char * src0,
+        device   const char * src1,
+        device         char * dst,
+        threadgroup    char * shmem [[threadgroup(0)]],
+        ushort3 tgpig[[threadgroup_position_in_grid]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort  tiisg[[thread_index_in_simdgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    constexpr short NW = N_SIMDWIDTH;
+
+    const short NSG = FC_solve_tri_nsg;
+    const short N   = FC_solve_tri_n;
+    const short K   = FC_solve_tri_k;
+    const short NP  = PAD2(N, NW);
+
+    const int32_t i03 = tgpig.z;
+    const int32_t i02 = tgpig.y;
+    const int32_t i01 = tgpig.x*NSG + sgitg;
+
+    threadgroup float * sh0 = (threadgroup float *) shmem;
+
+    device const float * src0_ptr = (device const float *)(src0 + i02 * args.nb02 + i03 * args.nb03) + sgitg*N;
+    device const float * src1_ptr = (device const float *)(src1 + i02 * args.nb12 + i03 * args.nb13) + i01;
+    device       float * dst_ptr  = (device       float *)(dst  + i02 * args.nb2  + i03 * args.nb3)  + i01;
+
+    for (short rr = 0; rr < N; rr += NSG) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        {
+            threadgroup float * sh0_cur = sh0 + sgitg*NP;
+
+            for (short t = 0; t*NW < N; ++t) {
+                const short idx = t*NW + tiisg;
+                sh0_cur[idx] = src0_ptr[idx];
+            }
+
+            src0_ptr += NSG*N;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (i01 >= args.ne10) {
+            continue;
+        }
+
+        for (short ir = 0; ir < NSG && rr + ir < N; ++ir) {
+            const short r = rr + ir;
+
+            threadgroup float * sh0_cur = sh0 + ir*NP;
+
+            float sum = 0.0f;
+
+            for (short t = 0; t*NW < r; ++t) {
+                const short idx = t*NW + tiisg;
+                sum += sh0_cur[idx] * dst_ptr[idx*K] * (idx < r);
+            }
+
+            sum = simd_sum(sum);
+
+            if (tiisg == 0) {
+                const float diag = sh0_cur[r];
+
+                dst_ptr[r*K] = (src1_ptr[r*K] - sum) / diag;
+            }
+        }
+    }
+}

ggml/src/ggml-metal/kernels/ssm.metal

@@ -0,0 +1,279 @@
+#include "common.h"
+
+// ref: ggml.c:ggml_compute_forward_ssm_conv_f32
+kernel void kernel_ssm_conv_f32_f32(
+        constant ggml_metal_kargs_ssm_conv & args,
+        device const  void * src0,
+        device const  void * src1,
+        device       float * dst,
+        uint3 tgpig[[threadgroup_position_in_grid]],
+        uint3 tpitg[[thread_position_in_threadgroup]],
+        uint3   ntg[[threads_per_threadgroup]]) {
+    const int64_t ir = tgpig.x;
+    const int64_t i2 = tgpig.y;
+    const int64_t i3 = tgpig.z;
+
+    const int64_t nc  = args.ne10;
+  //const int64_t ncs = args.ne00;
+  //const int64_t nr  = args.ne01;
+  //const int64_t n_t = args.ne1;
+  //const int64_t n_s = args.ne2;
+
+    device const float * s = (device const float *) ((device const char *) src0 + ir*args.nb01 + i2*args.nb00 + i3*args.nb02);
+    device const float * c = (device const float *) ((device const char *) src1 + ir*args.nb11);
+    device       float * x = (device       float *) ((device       char *) dst  + ir*args.nb0  + i2*args.nb1  + i3*args.nb2);
+
+    float sumf = 0.0f;
+
+    for (int64_t i0 = 0; i0 < nc; ++i0) {
+        sumf += s[i0] * c[i0];
+    }
+
+    x[0] = sumf;
+}
+
+kernel void kernel_ssm_conv_f32_f32_4(
+        constant ggml_metal_kargs_ssm_conv & args,
+        device const  void * src0,
+        device const  void * src1,
+        device       float * dst,
+        uint3 tgpig[[threadgroup_position_in_grid]],
+        uint3 tpitg[[thread_position_in_threadgroup]],
+        uint3   ntg[[threads_per_threadgroup]]) {
+    const int64_t ir = tgpig.x;
+    const int64_t i2 = tgpig.y;
+    const int64_t i3 = tgpig.z;
+
+    const int64_t nc  = args.ne10;
+  //const int64_t ncs = args.ne00;
+  //const int64_t nr  = args.ne01;
+  //const int64_t n_t = args.ne1;
+  //const int64_t n_s = args.ne2;
+
+    device const float4 * s = (device const float4 *) ((device const char *) src0 + ir*args.nb01 + i2*args.nb00 + i3*args.nb02);
+    device const float4 * c = (device const float4 *) ((device const char *) src1 + ir*args.nb11);
+    device       float  * x = (device       float  *) ((device       char *) dst  + ir*args.nb0  + i2*args.nb1  + i3*args.nb2);
+
+    float sumf = 0.0f;
+
+    for (int64_t i0 = 0; i0 < nc/4; ++i0) {
+        sumf += dot(s[i0], c[i0]);
+    }
+
+    x[0] = sumf;
+}
+
+constant short FC_ssm_conv_bs   [[function_constant(FC_SSM_CONV + 0)]];
+
+// Batched version: each threadgroup processes multiple tokens for better efficiency
+// Thread layout: each thread handles one token, threadgroup covers BATCH_SIZE tokens
+kernel void kernel_ssm_conv_f32_f32_batched(
+        constant ggml_metal_kargs_ssm_conv & args,
+        device const  void * src0,
+        device const  void * src1,
+        device       float * dst,
+        uint3 tgpig[[threadgroup_position_in_grid]],
+        uint3 tpitg[[thread_position_in_threadgroup]],
+        uint3   ntg[[threads_per_threadgroup]]) {
+    // tgpig.x = row index (ir)
+    // tgpig.y = batch of tokens (i2_base / BATCH_SIZE)
+    // tgpig.z = sequence index (i3)
+    // tpitg.x = thread within batch (0..BATCH_SIZE-1)
+    const short BATCH_SIZE = FC_ssm_conv_bs;
+
+    const int64_t ir      = tgpig.x;
+    const int64_t i2_base = tgpig.y * BATCH_SIZE;
+    const int64_t i3      = tgpig.z;
+    const int64_t i2_off  = tpitg.x;
+    const int64_t i2      = i2_base + i2_off;
+
+    const int64_t nc  = args.ne10;  // conv kernel size (typically 4)
+    const int64_t n_t = args.ne1;   // number of tokens
+
+    // Bounds check for partial batches at the end
+    if (i2 >= n_t) {
+        return;
+    }
+
+    // Load conv weights (shared across all tokens for this row)
+    device const float * c = (device const float *) ((device const char *) src1 + ir*args.nb11);
+
+    // Load source for this specific token
+    device const float * s = (device const float *) ((device const char *) src0 + ir*args.nb01 + i2*args.nb00 + i3*args.nb02);
+
+    // Output location for this token
+    device float * x = (device float *) ((device char *) dst + ir*args.nb0 + i2*args.nb1 + i3*args.nb2);
+
+    float sumf = 0.0f;
+    for (int64_t i0 = 0; i0 < nc; ++i0) {
+        sumf += s[i0] * c[i0];
+    }
+
+    x[0] = sumf;
+}
+
+kernel void kernel_ssm_conv_f32_f32_batched_4(
+        constant ggml_metal_kargs_ssm_conv & args,
+        device const  void * src0,
+        device const  void * src1,
+        device       float * dst,
+        uint3 tgpig[[threadgroup_position_in_grid]],
+        uint3 tpitg[[thread_position_in_threadgroup]],
+        uint3   ntg[[threads_per_threadgroup]]) {
+    // tgpig.x = row index (ir)
+    // tgpig.y = batch of tokens (i2_base / BATCH_SIZE)
+    // tgpig.z = sequence index (i3)
+    // tpitg.x = thread within batch (0..BATCH_SIZE-1)
+    const short BATCH_SIZE = FC_ssm_conv_bs;
+
+    const int64_t ir      = tgpig.x;
+    const int64_t i2_base = tgpig.y * BATCH_SIZE;
+    const int64_t i3      = tgpig.z;
+    const int64_t i2_off  = tpitg.x;
+    const int64_t i2      = i2_base + i2_off;
+
+    const int64_t nc  = args.ne10;  // conv kernel size (typically 4)
+    const int64_t n_t = args.ne1;   // number of tokens
+
+    // Bounds check for partial batches at the end
+    if (i2 >= n_t) {
+        return;
+    }
+
+    // Load conv weights (shared across all tokens for this row)
+    device const float4 * c = (device const float4 *) ((device const char *) src1 + ir*args.nb11);
+
+    // Load source for this specific token
+    device const float4 * s = (device const float4 *) ((device const char *) src0 + ir*args.nb01 + i2*args.nb00 + i3*args.nb02);
+
+    // Output location for this token
+    device float * x = (device float *) ((device char *) dst + ir*args.nb0 + i2*args.nb1 + i3*args.nb2);
+
+    float sumf = 0.0f;
+    for (int64_t i0 = 0; i0 < nc/4; ++i0) {
+        sumf += dot(s[i0], c[i0]);
+    }
+
+    x[0] = sumf;
+}
+
+// ref: ggml.c:ggml_compute_forward_ssm_scan_f32, Mamba-2 part
+// Optimized version: reduces redundant memory loads by having one thread load shared values
+kernel void kernel_ssm_scan_f32(
+        constant ggml_metal_kargs_ssm_scan & args,
+        device const void * src0,
+        device const void * src1,
+        device const void * src2,
+        device const void * src3,
+        device const void * src4,
+        device const void * src5,
+        device const void * src6,
+        device      float * dst,
+        threadgroup float * shared [[threadgroup(0)]],
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort  tiisg[[thread_index_in_simdgroup]],
+        ushort  sgptg[[simdgroups_per_threadgroup]],
+        uint3    tgpg[[threadgroups_per_grid]]) {
+    constexpr short NW = N_SIMDWIDTH;
+
+    // Shared memory layout:
+    // [0..sgptg*NW-1]: partial sums for reduction (existing)
+    // [sgptg*NW..sgptg*NW+sgptg-1]: pre-computed x_dt values for each token in batch
+    // [sgptg*NW+sgptg..sgptg*NW+2*sgptg-1]: pre-computed dA values for each token in batch
+    threadgroup float * shared_sums = shared;
+    threadgroup float * shared_x_dt = shared + sgptg * NW;
+    threadgroup float * shared_dA   = shared + sgptg * NW + sgptg;
+
+    shared_sums[tpitg.x] = 0.0f;
+
+    const int32_t i0 = tpitg.x;
+    const int32_t i1 = tgpig.x;
+    const int32_t ir = tgpig.y; // current head
+    const int32_t i3 = tgpig.z; // current seq
+
+    const int32_t nc  = args.d_state;
+    const int32_t nr  = args.d_inner;
+    const int32_t nh  = args.n_head;
+    const int32_t ng  = args.n_group;
+    const int32_t n_t = args.n_seq_tokens;
+
+    const int32_t s_off = args.s_off;
+
+    device const int32_t * ids = (device const int32_t *) src6;
+
+    device const float * s0_buff = (device const float *) ((device const char *) src0 + ir*args.nb02 + ids[i3]*args.nb03);
+    device       float * s_buff  = (device       float *) ((device       char *) dst  + ir*args.nb02 +      i3*args.nb03 + s_off);
+
+    const int32_t i = i0 + i1*nc;
+    const int32_t g = ir / (nh / ng); // repeat_interleave
+
+    float s0 = s0_buff[i];
+    float s  = 0.0f;
+
+    device const float * A = (device const float *) ((device const char *) src3 + ir*args.nb31); // {ne30, nh}
+
+    const float A0 = A[i0%args.ne30];
+
+    device const float * x  = (device const float *)((device const char *) src1 + i1*args.nb10  + ir*args.nb11 + i3*args.nb13); // {dim, nh, nt, ns}
+    device const float * dt = (device const float *)((device const char *) src2 + ir*args.nb20  + i3*args.nb22);                // {nh, nt, ns}
+    device const float * B  = (device const float *)((device const char *) src4 +  g*args.nb41  + i3*args.nb43);                // {d_state, ng, nt, ns}
+    device const float * C  = (device const float *)((device const char *) src5 +  g*args.nb51  + i3*args.nb53);                // {d_state, ng, nt, ns}
+
+    device float * y = dst + (i1 + ir*(nr) + i3*(n_t*nh*nr)); // {dim, nh, nt, ns}
+
+    for (int i2 = 0; i2 < n_t; i2 += sgptg) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Pre-compute x_dt and dA for this batch of tokens
+        // Only first sgptg threads do the loads and expensive math
+        if (i0 < sgptg && i2 + i0 < n_t) {
+            // ns12 and ns21 are element strides (nb12/nb10, nb21/nb20)
+            device const float * x_t  = x  + i0 * args.ns12;
+            device const float * dt_t = dt + i0 * args.ns21;
+
+            const float dt0  = dt_t[0];
+            const float dtsp = dt0 <= 20.0f ? log(1.0f + exp(dt0)) : dt0;
+            shared_x_dt[i0] = x_t[0] * dtsp;
+            shared_dA[i0]   = dtsp;  // Store dtsp, compute exp(dtsp * A0) per-thread since A0 varies
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        for (int t = 0; t < sgptg && i2 + t < n_t; t++) {
+            const float x_dt = shared_x_dt[t];
+            const float dA   = exp(shared_dA[t] * A0);
+
+            s = (s0 * dA) + (B[i0] * x_dt);
+
+            const float sumf = simd_sum(s * C[i0]);
+
+            if (tiisg == 0) {
+                shared_sums[t*NW + sgitg] = sumf;
+            }
+
+            // recurse
+            s0 = s;
+
+            B  += args.ns42;
+            C  += args.ns52;
+        }
+
+        // Advance pointers for next batch
+        x  += sgptg * args.ns12;
+        dt += sgptg * args.ns21;
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        const float sumf = simd_sum(shared_sums[sgitg*NW + tiisg]);
+
+        if (tiisg == 0 && i2 + sgitg < n_t) {
+            y[sgitg*nh*nr] = sumf;
+        }
+
+        y += sgptg*nh*nr;
+    }
+
+    s_buff[i] = s;
+}

ggml/src/ggml-metal/kernels/tri.metal

@@ -0,0 +1,69 @@
+#include "common.h"
+
+template<uint32_t ttype>
+bool _ggml_vec_tri_cmp(const int i, const int r);
+
+template<>
+bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_LOWER */ 3>(const int i, const int r) {
+    return i < r;
+}
+
+template<>
+bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_LOWER_DIAG */ 2>(const int i, const int r) {
+    return i <= r;
+}
+
+template<>
+bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_UPPER */ 1>(const int i, const int r) {
+    return i > r;
+}
+
+template<>
+bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_UPPER_DIAG */ 0>(const int i, const int r) {
+    return i >= r;
+}
+
+template<typename T, int ttype>
+kernel void kernel_tri(
+        constant ggml_metal_kargs_tri & args,
+        device const char * src0,
+        device const char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    const int i3 = tgpig.z;
+    const int i2 = tgpig.y;
+    const int i1 = tgpig.x;
+
+    if (i3 >= args.ne03 || i2 >= args.ne02 || i1 >= args.ne01) {
+        return;
+    }
+
+    device const T * src_row = (device const T *) ((device const char *) src0 + i1*args.nb01 + i2*args.nb02 + i3*args.nb03);
+    device       T * dst_row = (device       T *) ((device       char *) dst  + i1*args.nb1  + i2*args.nb2  + i3*args.nb3);
+
+    // Each thread is a single element of the row if ne00 < max threads per
+    // threadgroup, so this will loop once for each index that this thread is
+    // responsible for
+    for (int64_t i0 = tpitg.x; i0 < args.ne00; i0 += ntg.x) {
+        // Use the comparison as a mask for branchless
+        dst_row[i0] = static_cast<T>(_ggml_vec_tri_cmp<ttype>(i0, i1)) * src_row[i0];
+    }
+}
+
+typedef decltype(kernel_tri<float, 0>) kernel_tri_t;
+
+template [[host_name("kernel_tri_f32_0")]] kernel kernel_tri_t kernel_tri<float, 0>;
+template [[host_name("kernel_tri_f32_1")]] kernel kernel_tri_t kernel_tri<float, 1>;
+template [[host_name("kernel_tri_f32_2")]] kernel kernel_tri_t kernel_tri<float, 2>;
+template [[host_name("kernel_tri_f32_3")]] kernel kernel_tri_t kernel_tri<float, 3>;
+template [[host_name("kernel_tri_f16_0")]] kernel kernel_tri_t kernel_tri<half, 0>;
+template [[host_name("kernel_tri_f16_1")]] kernel kernel_tri_t kernel_tri<half, 1>;
+template [[host_name("kernel_tri_f16_2")]] kernel kernel_tri_t kernel_tri<half, 2>;
+template [[host_name("kernel_tri_f16_3")]] kernel kernel_tri_t kernel_tri<half, 3>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_tri_bf16_0")]] kernel kernel_tri_t kernel_tri<bfloat, 0>;
+template [[host_name("kernel_tri_bf16_1")]] kernel kernel_tri_t kernel_tri<bfloat, 1>;
+template [[host_name("kernel_tri_bf16_2")]] kernel kernel_tri_t kernel_tri<bfloat, 2>;
+template [[host_name("kernel_tri_bf16_3")]] kernel kernel_tri_t kernel_tri<bfloat, 3>;
+#endif

ggml/src/ggml-metal/kernels/unary.metal

@@ -0,0 +1,360 @@
+#include "common.h"
+
+constant short FC_unary_op [[function_constant(FC_UNARY + 0)]];
+constant bool  FC_unary_cnt[[function_constant(FC_UNARY + 1)]];
+
+template <typename T0, typename T, typename TC>
+kernel void kernel_unary_impl(
+        constant ggml_metal_kargs_unary & args,
+        device const char * src0,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+#define FC_OP  FC_unary_op
+#define FC_CNT FC_unary_cnt
+
+    device const T0 * src0_ptr;
+    device       T  * dst_ptr;
+
+    int i0;
+
+    if (FC_CNT) {
+        i0 = tgpig.x;
+
+        src0_ptr = (device const T0 *) (src0);
+        dst_ptr  = (device       T  *) (dst);
+    } else {
+        const int i03 = tgpig.z;
+        const int i02 = tgpig.y;
+        const int k0  = tgpig.x/args.ne01;
+        const int i01 = tgpig.x - k0*args.ne01;
+
+        i0 = k0*ntg.x + tpitg.x;
+
+        src0_ptr = (device const T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
+        dst_ptr  = (device       T  *) (dst  + i03*args.nb3  + i02*args.nb2  + i01*args.nb1 );
+    }
+
+    {
+        //threadgroup_barrier(mem_flags::mem_none);
+
+        if (!FC_CNT) {
+            if (i0 >= args.ne0) {
+                return;
+            }
+        }
+
+        const TC x = (TC) src0_ptr[i0];
+
+        if (FC_OP == OP_UNARY_NUM_SCALE) {
+            dst_ptr[i0] = (T) (args.scale * x + args.bias);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_FILL) {
+            dst_ptr[i0] = (T) args.val;
+        }
+
+        if (FC_OP == OP_UNARY_NUM_CLAMP) {
+            dst_ptr[i0] = (T) clamp(x, args.min, args.max);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_SQR) {
+            dst_ptr[i0] = (T) (x * x);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_SQRT) {
+            dst_ptr[i0] = (T) sqrt(x);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_SIN) {
+            dst_ptr[i0] = (T) sin(x);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_COS) {
+            dst_ptr[i0] = (T) cos(x);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_LOG) {
+            dst_ptr[i0] = (T) log(x);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_LEAKY_RELU) {
+            dst_ptr[i0] = (T) (TC(x > 0)*x + TC(x <= 0)*(x * args.slope));
+        }
+
+        if (FC_OP == OP_UNARY_NUM_TANH) {
+            dst_ptr[i0] = (T) precise::tanh(x);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_RELU) {
+            dst_ptr[i0] = (T) fmax(0, x);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_SIGMOID) {
+            dst_ptr[i0] = (T) (1 / (1 + exp(-x)));
+        }
+
+        if (FC_OP == OP_UNARY_NUM_GELU) {
+            dst_ptr[i0] = (T) (0.5*x*(1 + precise::tanh(SQRT_2_OVER_PI*x*(1 + GELU_COEF_A*x*x))));
+        }
+
+        if (FC_OP == OP_UNARY_NUM_GELU_ERF) {
+            dst_ptr[i0] = (T) (0.5*x*(1 + erf_approx(SQRT_2_INV*x)));
+        }
+
+        if (FC_OP == OP_UNARY_NUM_GELU_QUICK) {
+            dst_ptr[i0] = (T) (x * (1/(1 + exp(GELU_QUICK_COEF*x))));
+        }
+
+        if (FC_OP == OP_UNARY_NUM_SILU) {
+            dst_ptr[i0] = (T) (x / (1 + exp(-x)));
+        }
+
+        if (FC_OP == OP_UNARY_NUM_ELU) {
+            dst_ptr[i0] = (T) elu_approx(x);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_NEG) {
+            dst_ptr[i0] = (T) -x;
+        }
+
+        if (FC_OP == OP_UNARY_NUM_ABS) {
+            dst_ptr[i0] = (T) fabs(x);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_SGN) {
+            dst_ptr[i0] = T(x > 0) - T(x < 0);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_STEP) {
+            dst_ptr[i0] = T(x > 0);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_HARDSWISH) {
+            dst_ptr[i0] = (T) (x * fmax(0, fmin(1, x/6 + 0.5)));
+        }
+
+        if (FC_OP == OP_UNARY_NUM_HARDSIGMOID) {
+            dst_ptr[i0] = (T) fmax(0, fmin(1, x/6 + 0.5));
+        }
+
+        if (FC_OP == OP_UNARY_NUM_EXP) {
+            dst_ptr[i0] = (T) exp(x);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_SOFTPLUS) {
+            dst_ptr[i0] = (T) select(log(1 + exp(x)), x, x > 20);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_EXPM1) {
+            // TODO: precise implementation
+            dst_ptr[i0] = (T) (exp(x) - 1);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_FLOOR) {
+            dst_ptr[i0] = (T) floor(x);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_CEIL) {
+            dst_ptr[i0] = (T) ceil(x);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_ROUND) {
+            dst_ptr[i0] = (T) round(x);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_TRUNC) {
+            dst_ptr[i0] = (T) trunc(x);
+        }
+
+        if (FC_OP == OP_UNARY_NUM_XIELU) {
+            const TC xi      = x;
+            const TC gate    = TC(xi > TC(0.0f));
+            const TC clamped = fmin(xi, TC(args.val));
+            const TC y_pos   = TC(args.scale) * xi * xi + TC(args.bias) * xi;
+            const TC y_neg   = (exp(clamped) - TC(1.0f) - xi) * TC(args.slope) + TC(args.bias) * xi;
+            dst_ptr[i0] = (T) (gate * y_pos + (TC(1.0f) - gate) * y_neg);
+        }
+    }
+
+#undef FC_OP
+#undef FC_CNT
+}
+
+typedef decltype(kernel_unary_impl<float, float, float>) kernel_unary_t;
+
+template [[host_name("kernel_unary_f32_f32")]]   kernel kernel_unary_t kernel_unary_impl<float,  float,  float>;
+template [[host_name("kernel_unary_f32_f32_4")]] kernel kernel_unary_t kernel_unary_impl<float4, float4, float4>;
+template [[host_name("kernel_unary_f16_f16")]]   kernel kernel_unary_t kernel_unary_impl<half,   half,   float>;
+template [[host_name("kernel_unary_f16_f16_4")]] kernel kernel_unary_t kernel_unary_impl<half4,  half4,  float4>;
+
+template<typename T>
+kernel void kernel_reglu(
+        constant ggml_metal_kargs_glu & args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        uint tgpig[[threadgroup_position_in_grid]],
+        uint tpitg[[thread_position_in_threadgroup]],
+        uint   ntg[[threads_per_threadgroup]]) {
+    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
+    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
+    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
+
+    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+        const float x0 = src0_row[i0];
+        const float x1 = src1_row[i0];
+
+        dst_row[i0] = (T)(x0*x1*(x0 > 0.0f));
+    }
+}
+
+typedef decltype(kernel_reglu<float>) kernel_reglu_t;
+
+template [[host_name("kernel_reglu_f32")]] kernel kernel_reglu_t kernel_reglu<float>;
+template [[host_name("kernel_reglu_f16")]] kernel kernel_reglu_t kernel_reglu<half>;
+
+template<typename T>
+kernel void kernel_geglu(
+        constant ggml_metal_kargs_glu & args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        uint tgpig[[threadgroup_position_in_grid]],
+        uint tpitg[[thread_position_in_threadgroup]],
+        uint   ntg[[threads_per_threadgroup]]) {
+    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
+    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
+    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
+
+    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+        const float x0 = src0_row[i0];
+        const float x1 = src1_row[i0];
+
+        const float gelu = 0.5f*x0*(1.0f + precise::tanh(SQRT_2_OVER_PI*x0*(1.0f + GELU_COEF_A*x0*x0)));
+
+        dst_row[i0] = (T)(gelu*x1);
+    }
+}
+
+typedef decltype(kernel_geglu<float>) kernel_geglu_t;
+
+template [[host_name("kernel_geglu_f32")]] kernel kernel_geglu_t kernel_geglu<float>;
+template [[host_name("kernel_geglu_f16")]] kernel kernel_geglu_t kernel_geglu<half>;
+
+template<typename T>
+kernel void kernel_swiglu(
+        constant ggml_metal_kargs_glu & args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        uint tgpig[[threadgroup_position_in_grid]],
+        uint tpitg[[thread_position_in_threadgroup]],
+        uint   ntg[[threads_per_threadgroup]]) {
+    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
+    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
+    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
+
+    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+        const float x0 = src0_row[i0];
+        const float x1 = src1_row[i0];
+
+        const float silu = x0 / (1.0f + exp(-x0));
+
+        dst_row[i0] = (T)(silu*x1);
+    }
+}
+
+typedef decltype(kernel_swiglu<float>) kernel_swiglu_t;
+
+template [[host_name("kernel_swiglu_f32")]] kernel kernel_swiglu_t kernel_swiglu<float>;
+template [[host_name("kernel_swiglu_f16")]] kernel kernel_swiglu_t kernel_swiglu<half>;
+
+template<typename T>
+kernel void kernel_swiglu_oai(
+        constant ggml_metal_kargs_glu & args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        uint tgpig[[threadgroup_position_in_grid]],
+        uint tpitg[[thread_position_in_threadgroup]],
+        uint   ntg[[threads_per_threadgroup]]) {
+    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
+    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
+    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
+
+    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+        float x0 = src0_row[i0];
+        float x1 = src1_row[i0];
+
+        x0 = min(x0, args.limit);
+        x1 = max(min(x1, args.limit), -args.limit);
+
+        float out_glu = x0 / (1.0f + exp(-x0 * args.alpha));
+        out_glu = out_glu * (1.0f + x1);
+
+        dst_row[i0] = (T)out_glu;
+    }
+}
+
+typedef decltype(kernel_swiglu_oai<float>) kernel_swiglu_oai_t;
+
+template [[host_name("kernel_swiglu_oai_f32")]] kernel kernel_swiglu_oai_t kernel_swiglu_oai<float>;
+template [[host_name("kernel_swiglu_oai_f16")]] kernel kernel_swiglu_oai_t kernel_swiglu_oai<half>;
+
+template<typename T>
+kernel void kernel_geglu_erf(
+        constant ggml_metal_kargs_glu & args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        uint tgpig[[threadgroup_position_in_grid]],
+        uint tpitg[[thread_position_in_threadgroup]],
+        uint   ntg[[threads_per_threadgroup]]) {
+    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
+    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
+    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
+
+    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+        const float x0 = src0_row[i0];
+        const float x1 = src1_row[i0];
+
+        const float gelu_erf = 0.5f*x0*(1.0f+erf_approx<float>(x0*SQRT_2_INV));
+
+        dst_row[i0] = (T)(gelu_erf*x1);
+    }
+}
+
+typedef decltype(kernel_geglu_erf<float>) kernel_geglu_erf_t;
+
+template [[host_name("kernel_geglu_erf_f32")]] kernel kernel_geglu_erf_t kernel_geglu_erf<float>;
+template [[host_name("kernel_geglu_erf_f16")]] kernel kernel_geglu_erf_t kernel_geglu_erf<half>;
+
+template<typename T>
+kernel void kernel_geglu_quick(
+        constant ggml_metal_kargs_glu & args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        uint tgpig[[threadgroup_position_in_grid]],
+        uint tpitg[[thread_position_in_threadgroup]],
+        uint   ntg[[threads_per_threadgroup]]) {
+    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
+    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
+    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
+
+    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+        const float x0 = src0_row[i0];
+        const float x1 = src1_row[i0];
+
+        const float gelu_quick = x0*(1.0f/(1.0f+exp(GELU_QUICK_COEF*x0)));
+
+        dst_row[i0] = (T)(gelu_quick*x1);
+    }
+}
+
+typedef decltype(kernel_geglu_quick<float>) kernel_geglu_quick_t;
+
+template [[host_name("kernel_geglu_quick_f32")]] kernel kernel_geglu_quick_t kernel_geglu_quick<float>;
+template [[host_name("kernel_geglu_quick_f16")]] kernel kernel_geglu_quick_t kernel_geglu_quick<half>;

ggml/src/ggml-metal/kernels/upscale.metal

@@ -0,0 +1,179 @@
+#include "common.h"
+
+constant bool FC_upscale_aa [[function_constant(FC_UPSCALE + 0)]];
+
+kernel void kernel_upscale_nearest_f32(
+    constant ggml_metal_kargs_upscale & args,
+    device  const char * src0,
+    device        char * dst,
+    uint3 tgpig[[threadgroup_position_in_grid]],
+    uint3 tpitg[[thread_position_in_threadgroup]],
+    uint3   ntg[[threads_per_threadgroup]]) {
+
+    const int64_t i3 = tgpig.z;
+    const int64_t i2 = tgpig.y;
+    const int64_t i1 = tgpig.x;
+
+    const int64_t i03 = i3/args.sf3;
+    const int64_t i02 = i2/args.sf2;
+    const int64_t i01 = i1/args.sf1;
+
+    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
+        const int64_t i00 = i0/args.sf0;
+
+        device const float * src0_ptr = (device const float *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + i00*args.nb00);
+        device       float * dst_ptr  = (device       float *) (dst  +  i3*args.nb3  +  i2*args.nb2  +  i1*args.nb1  +  i0*args.nb0);
+
+        dst_ptr[0] = src0_ptr[0];
+    }
+}
+
+static inline float bilinear_tri(float x) {
+    return MAX(0.0f, 1.0f - fabs(x));
+}
+
+kernel void kernel_upscale_bilinear_f32(
+    constant ggml_metal_kargs_upscale & args,
+    device  const char * src0,
+    device        char * dst,
+    uint3 tgpig[[threadgroup_position_in_grid]],
+    uint3 tpitg[[thread_position_in_threadgroup]],
+    uint3   ntg[[threads_per_threadgroup]]) {
+
+    const int64_t i3 = tgpig.z;
+    const int64_t i2 = tgpig.y;
+    const int64_t i1 = tgpig.x;
+
+    const int64_t i03 = i3 / args.sf3;
+    const int64_t i02 = i2 / args.sf2;
+
+    const float   f01  = ((float)i1 + args.poffs) / args.sf1 - args.poffs;
+    const int64_t i01  = MAX(0, MIN(args.ne01 - 1, (int64_t)floor(f01)));
+    const int64_t i01p = MAX(0, MIN(args.ne01 - 1, i01 + 1));
+    const float   fd1  = MAX(0.0f, MIN(1.0f, f01 - (float)i01));
+
+    src0 += i03*args.nb03 + i02*args.nb02;
+
+    device float * dst_ptr = (device float *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1);
+
+    if (FC_upscale_aa) {
+        const float support0  = MAX(1.0f, 1.0f / args.sf0);
+        const float invscale0 = 1.0f / support0;
+        const float support1  = MAX(1.0f, 1.0f / args.sf1);
+        const float invscale1 = 1.0f / support1;
+
+        for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
+            const float f00 = ((float)i0 + args.poffs) / args.sf0 - args.poffs;
+
+            int64_t x_min = MAX((int64_t)0, (int64_t)floor(f00 - support0 + args.poffs));
+            int64_t x_max = MIN(args.ne00,  (int64_t)ceil (f00 + support0 + args.poffs));
+
+            int64_t y_min = MAX((int64_t)0, (int64_t)floor(f01 - support1 + args.poffs));
+            int64_t y_max = MIN(args.ne01,  (int64_t)ceil (f01 + support1 + args.poffs));
+
+            float sum = 0.0f;
+            float wsum = 0.0f;
+
+            for (int64_t sy = y_min; sy < y_max; ++sy) {
+                const float wy = MAX(0.0f, 1.0f - fabs((float)sy - f01) * invscale1);
+                for (int64_t sx = x_min; sx < x_max; ++sx) {
+                    const float wx = MAX(0.0f, 1.0f - fabs((float)sx - f00) * invscale0);
+                    const float w  = wx * wy;
+                    device const float * src_ptr = (device const float *)(src0 + sy*args.nb01 + sx*args.nb00);
+                    sum  += (*src_ptr) * w;
+                    wsum += w;
+                }
+            }
+
+            const float v = (wsum > 0.0f) ? (sum / wsum) : 0.0f;
+            dst_ptr[i0] = v;
+        }
+    } else {
+        for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
+            const float   f00  = ((float)i0 + args.poffs) / args.sf0 - args.poffs;
+            const int64_t i00  = MAX(0, MIN(args.ne00 - 1, (int64_t)floor(f00)));
+            const int64_t i00p = MAX(0, MIN(args.ne00 - 1, i00 + 1));
+            const float   fd0  = MAX(0.0f, MIN(1.0f, f00 - (float)i00));
+
+            device const float * src00 = (device const float *)(src0 + i01*args.nb01  + i00*args.nb00);
+            device const float * src10 = (device const float *)(src0 + i01*args.nb01  + i00p*args.nb00);
+            device const float * src01 = (device const float *)(src0 + i01p*args.nb01 + i00*args.nb00);
+            device const float * src11 = (device const float *)(src0 + i01p*args.nb01 + i00p*args.nb00);
+
+            const float v =
+                (*src00) * (1.0f - fd0) * (1.0f - fd1) +
+                (*src10) * fd0          * (1.0f - fd1) +
+                (*src01) * (1.0f - fd0) * fd1 +
+                (*src11) * fd0          * fd1;
+
+            dst_ptr[i0] = v;
+        }
+    }
+}
+
+static inline float bicubic_weight1(float x) {
+    const float a = -0.75f;
+    return ((a + 2) * x - (a + 3)) * x * x + 1;
+}
+
+static inline float bicubic_weight2(float x) {
+    const float a = -0.75f;
+    return ((a * x - 5 * a) * x + 8 * a) * x - 4 * a;
+}
+
+kernel void kernel_upscale_bicubic_f32(
+    constant ggml_metal_kargs_upscale & args,
+    device  const char * src0,
+    device        char * dst,
+    uint3 tgpig[[threadgroup_position_in_grid]],
+    uint3 tpitg[[thread_position_in_threadgroup]],
+    uint3   ntg[[threads_per_threadgroup]]) {
+
+    const int64_t i3 = tgpig.z;
+    const int64_t i2 = tgpig.y;
+    const int64_t i1 = tgpig.x;
+
+    const int64_t i03 = i3 / args.sf3;
+    const int64_t i02 = i2 / args.sf2;
+
+    const float   f01 = ((float)i1 + args.poffs) / args.sf1 - args.poffs;
+    const int64_t i01 = (int64_t)floor(f01);
+    const float   fd1 = f01 - (float)i01;
+
+    const float w_y0 = bicubic_weight2(fd1 + 1.0f);
+    const float w_y1 = bicubic_weight1(fd1);
+    const float w_y2 = bicubic_weight1(1.0f - fd1);
+    const float w_y3 = bicubic_weight2(2.0f - fd1);
+
+    const device const char * src_slice = src0 + i03 * args.nb03 + i02 * args.nb02;
+
+    device float * dst_ptr = (device float *)(dst + i3 * args.nb3 + i2 * args.nb2 + i1 * args.nb1);
+
+    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
+        const float   f00 = ((float)i0 + args.poffs) / args.sf0 - args.poffs;
+        const int64_t i00 = (int64_t)floor(f00);
+        const float   fd0 = f00 - (float)i00;
+
+        const float w_x0 = bicubic_weight2(fd0 + 1.0f);
+        const float w_x1 = bicubic_weight1(fd0);
+        const float w_x2 = bicubic_weight1(1.0f - fd0);
+        const float w_x3 = bicubic_weight2(2.0f - fd0);
+
+        float sum = 0.0f;
+
+        for (int dy = -1; dy <= 2; ++dy) {
+            const int64_t iy = MAX(0, MIN(args.ne01 - 1, i01 + dy));
+            const float wy = (dy == -1) ? w_y0 : (dy == 0) ? w_y1 : (dy == 1) ? w_y2 : w_y3;
+
+            for (int dx = -1; dx <= 2; ++dx) {
+                const int64_t ix = MAX(0, MIN(args.ne00 - 1, i00 + dx));
+                const float wx = (dx == -1) ? w_x0 : (dx == 0) ? w_x1 : (dx == 1) ? w_x2 : w_x3;
+
+                device const float * src_ptr = (device const float *)(src_slice + iy * args.nb01 + ix * args.nb00);
+                sum += (*src_ptr) * wx * wy;
+            }
+        }
+
+        dst_ptr[i0] = sum;
+    }
+}

ggml/src/ggml-metal/kernels/wkv.metal

@@ -0,0 +1,179 @@
+#include "common.h"
+
+kernel void kernel_rwkv_wkv6_f32(
+    device const float * k,
+    device const float * v,
+    device const float * r,
+    device const float * tf,
+    device const float * td,
+    device const float * state_in,
+    device       float * dst,
+    constant    uint & B,
+    constant    uint & T,
+    constant    uint & C,
+    constant    uint & H,
+    uint3 tgpig[[threadgroup_position_in_grid]],
+    uint3 tpitg[[thread_position_in_threadgroup]],
+    uint3   ntg[[threads_per_threadgroup]])  {
+
+    const uint head_size = 64; // TODO: support head_size = 128
+    const uint batch_id = tgpig.x / H;
+    const uint head_id = tgpig.x % H;
+    const uint tid = tpitg.x;
+
+    if (batch_id >= B || head_id >= H) {
+        return;
+    }
+
+    const uint state_size = C * head_size;
+    const uint n_seq_tokens = T / B;
+
+    threadgroup float _k[head_size];
+    threadgroup float _r[head_size];
+    threadgroup float _tf[head_size];
+    threadgroup float _td[head_size];
+
+    float state[head_size];
+
+    for (uint i = 0; i < head_size; i++) {
+        state[i] = state_in[batch_id * state_size + head_id * head_size * head_size
+                          + i * head_size + tid];
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    _tf[tid] = tf[head_id * head_size + tid];
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid;
+    const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid;
+
+    for (uint t = start_t; t < end_t; t += C) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        _k[tid] = k[t];
+        _r[tid] = r[t];
+        _td[tid] = td[t];
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        const float v_val = v[t];
+        float y = 0.0;
+
+        for (uint j = 0; j < head_size; j += 4) {
+            float4 k_vec = float4(_k[j], _k[j+1], _k[j+2], _k[j+3]);
+            float4 r_vec = float4(_r[j], _r[j+1], _r[j+2], _r[j+3]);
+            float4 tf_vec = float4(_tf[j], _tf[j+1], _tf[j+2], _tf[j+3]);
+            float4 td_vec = float4(_td[j], _td[j+1], _td[j+2], _td[j+3]);
+            float4 s_vec = float4(state[j], state[j+1], state[j+2], state[j+3]);
+
+            float4 kv = k_vec * v_val;
+
+            float4 temp = tf_vec * kv + s_vec;
+            y += dot(r_vec, temp);
+
+            s_vec = s_vec * td_vec + kv;
+            state[j]   = s_vec[0];
+            state[j+1] = s_vec[1];
+            state[j+2] = s_vec[2];
+            state[j+3] = s_vec[3];
+        }
+
+        dst[t] = y;
+    }
+
+    for (uint i = 0; i < head_size; i++) {
+        dst[T * C + batch_id * state_size + head_id * head_size * head_size
+            + i * head_size + tid] = state[i];
+    }
+}
+
+kernel void kernel_rwkv_wkv7_f32(
+    device const float * r,
+    device const float * w,
+    device const float * k,
+    device const float * v,
+    device const float * a,
+    device const float * b,
+    device const float * state_in,
+    device       float * dst,
+    constant    uint & B,
+    constant    uint & T,
+    constant    uint & C,
+    constant    uint & H,
+    uint3 tgpig[[threadgroup_position_in_grid]],
+    uint3 tpitg[[thread_position_in_threadgroup]],
+    uint3   ntg[[threads_per_threadgroup]])  {
+
+    const uint head_size = 64; // TODO: support head_size = 128
+    const uint batch_id = tgpig.x / H;
+    const uint head_id = tgpig.x % H;
+    const uint tid = tpitg.x;
+
+    if (batch_id >= B || head_id >= H) {
+        return;
+    }
+
+    const uint state_size = C * head_size;
+    const uint n_seq_tokens = T / B;
+
+    threadgroup float _r[head_size];
+    threadgroup float _w[head_size];
+    threadgroup float _k[head_size];
+    threadgroup float _a[head_size];
+    threadgroup float _b[head_size];
+
+    float state[head_size];
+
+    for (uint i = 0; i < head_size; i++) {
+        state[i] = state_in[batch_id * state_size + head_id * head_size * head_size
+                          + tid * head_size + i];
+    }
+
+    const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid;
+    const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid;
+
+    for (uint t = start_t; t < end_t; t += C) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        _r[tid] = r[t];
+        _w[tid] = w[t];
+        _k[tid] = k[t];
+        _a[tid] = a[t];
+        _b[tid] = b[t];
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        const float v_val = v[t];
+        float y = 0.0, sa = 0.0;
+
+        float4 sa_vec(0.0);
+
+        for (uint j = 0; j < head_size; j += 4) {
+            float4 a_vec = float4(_a[j], _a[j+1], _a[j+2], _a[j+3]);
+            float4 s_vec = float4(state[j], state[j+1], state[j+2], state[j+3]);
+            sa_vec += a_vec * s_vec;
+        }
+        sa = sa_vec[0] + sa_vec[1] + sa_vec[2] + sa_vec[3];
+
+        for (uint j = 0; j < head_size; j += 4) {
+            float4 r_vec = float4(_r[j], _r[j+1], _r[j+2], _r[j+3]);
+            float4 w_vec = float4(_w[j], _w[j+1], _w[j+2], _w[j+3]);
+            float4 k_vec = float4(_k[j], _k[j+1], _k[j+2], _k[j+3]);
+            float4 b_vec = float4(_b[j], _b[j+1], _b[j+2], _b[j+3]);
+            float4 s_vec = float4(state[j], state[j+1], state[j+2], state[j+3]);
+
+            float4 kv = k_vec * v_val;
+
+            s_vec = s_vec * w_vec + kv + sa * b_vec;
+            y += dot(s_vec, r_vec);
+
+            state[j]   = s_vec[0];
+            state[j+1] = s_vec[1];
+            state[j+2] = s_vec[2];
+            state[j+3] = s_vec[3];
+        }
+
+        dst[t] = y;
+    }
+
+    for (uint i = 0; i < head_size; i++) {
+        dst[T * C + batch_id * state_size + head_id * head_size * head_size
+            + tid * head_size + i] = state[i];
+    }
+}

ggml/src/ggml-opencl/CMakeLists.txt

@@ -114,9 +114,7 @@ set(GGML_OPENCL_KERNELS
     mul_mv_id_mxfp4_f32
     mul_mv_id_mxfp4_f32_flat
     gemm_moe_q4_0_f32_ns
-    gemm_moe_q4_0_q8_1_dp4a
     gemv_moe_q4_0_f32_ns
-    gemm_moe_q8_0_f32_ns
     gemm_moe_q4_1_f32_ns
     gemv_moe_q4_1_f32_ns
     gemm_moe_q5_0_f32_ns
@@ -124,18 +122,6 @@ set(GGML_OPENCL_KERNELS
     gemm_moe_q5_1_f32_ns
     gemv_moe_q5_1_f32_ns
     gemm_moe_q4_k_f32_ns
-    gemm_moe_q4_k_q8_1_dp4a
-    gemm_moe_q6_k_q8_1_dp4a
-    gemm_moe_q8_1_dp4a
-    moe_reorder_quant_a_q8_1
-    gemm_noshuffle_q4_k_q8_1_dp4a
-    gemm_noshuffle_q5_k_q8_1_dp4a
-    gemm_noshuffle_q6_k_q8_1_dp4a
-    gemm_noshuffle_q8_0_q8_1_dp4a
-    gemm_noshuffle_q5_0_q8_1_dp4a
-    gemm_noshuffle_iq4_nl_q8_1_dp4a
-    gemm_noshuffle_q4_0_q8_1_dp4a
-    quant_a_q8_1
     gemv_moe_q4_k_f32_ns
     gemm_moe_q5_k_f32_ns
     gemv_moe_q5_k_f32_ns
@@ -144,10 +130,8 @@ set(GGML_OPENCL_KERNELS
     gemm_moe_mxfp4_f32
     gemv_moe_mxfp4_f32
     gemm_moe_mxfp4_f32_ns
-    gemm_moe_mxfp4_q8_1_dp4a
     gemv_moe_mxfp4_f32_ns
     moe_reorder_b
-    moe_combine
     moe_sort_by_expert
     mul_mm_f32_f32_l4_lm
     mul_mm_f16_f32_l4_lm

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

@@ -2372,121 +2372,3 @@ kernel void kernel_restore_block_iq4_nl_noshuffle(
         b->qs[2*i + 1] = convert_uchar(((x0 & mask_F0) >> 4) | (x1 & mask_F0));
     }
 }
-
-// ---------------------------------------------------------------------------
-// kernel_moe_expand_scale_q8_0
-//
-// Expand the q8_0 per-32-block scale d (one half/block, [expert][row][block]) into
-// the UNIFORM scale[16] format the generic dp4a MoE GEMM (kernel_gemm_moe_q8_1_dp4a,
-// MOE_QT=80) consumes: 16 f16 per 256-superblock (per-16-element segment), where the
-// two segments of each 32-block share the block's d. q8_0 is symmetric -> no min
-// buffer (the GEMM runs with has_min=0). The int8 weight codes are reused verbatim
-// from the existing flat q8_0 weight buffer (extra0_q8_0->q), so only the scale is
-// rebuilt here. One work-item per (row, superblock, expert).
-// ---------------------------------------------------------------------------
-kernel void kernel_moe_expand_scale_q8_0(
-    global const half * src_d,      // [expert][row][block], one scale per 32-block
-    global       half * dst_scale,  // [expert][row][block][2] (FLAT per-32-block)
-    int ne00,
-    int ne01
-) {
-    int row = get_global_id(0);
-    int blk = get_global_id(1);   // 32-block index along K
-    int e   = get_global_id(2);
-    if (row >= ne01) { return; }
-
-    long nb = ne00 / 32;          // 32-blocks per row (K only needs % 32 == 0)
-    half d  = src_d[((long)e*ne01 + row)*nb + blk];
-    long b  = (((long)e*ne01 + row)*nb + blk) * 2;
-    dst_scale[b + 0] = d;
-    dst_scale[b + 1] = d;
-}
-
-// ---------------------------------------------------------------------------
-// kernel_moe_expand_scale_q5_0
-//
-// q5_0 = symmetric, value = d*(code-16), code = nibble | (hi<<4) in 0..31. The
-// generic dp4a MoE GEMM keeps the unsigned code and centers via the min term:
-//   scale*dp4a(code,a) - min*sum(a),  scale = d,  min = d*16.
-// Reads the existing q5_0 d ([expert][block][row], one half/32-block, from the
-// trans4 convert) and writes the FLAT per-32-block uniform scale[2]/min[1] in
-// [expert][row][block] order (a transpose). One work-item per (row, block, expert).
-// ---------------------------------------------------------------------------
-kernel void kernel_moe_expand_scale_q5_0(
-    global const half * src_d,      // [expert][block][row]
-    global       half * dst_scale,  // [expert][row][block][2]
-    global       half * dst_min,    // [expert][row][block]
-    int ne00,
-    int ne01
-) {
-    int row = get_global_id(0);
-    int blk = get_global_id(1);
-    int e   = get_global_id(2);
-    if (row >= ne01) { return; }
-
-    long nb = ne00 / 32;
-    half d  = src_d[(long)e*nb*ne01 + (long)blk*ne01 + row];   // [expert][block][row]
-    long sb = (((long)e*ne01 + row)*nb + blk) * 2;
-    long mb = ((long)e*ne01 + row)*nb + blk;
-    dst_scale[sb + 0] = d;
-    dst_scale[sb + 1] = d;
-    dst_min[mb] = (half)((float)d * 16.0f);
-}
-
-// ---------------------------------------------------------------------------
-// kernel_moe_expand_scale_q5_K
-//
-// q5_K value = d*sv*code + (-dm*mn), with the 6-bit packed per-sub-block scale sv
-// and min mn (8 sub-blocks of 32 per 256-superblock, decoded by get_scale_min_k4
-// from the 12-byte s[]). The generic dp4a MoE GEMM (kernel_gemm_moe_q8_1_dp4a,
-// MOE_QT=5) keeps the unsigned 5-bit code and applies scale/min via the uniform
-// per-32-block buffers:
-//   acc += sc0*a_d*raw1 + sc1*a_d*raw2 - mn_u*a_s,
-//   sc0 = sc1 = d*sv (both per-16 segments of a 32-block share the sub-block scale),
-//   mn_u = dm*mn (positive; the GEMM subtracts it -> the -dm*mn min term).
-// q5_K's q_img (low nibbles) + qh (hi-bit plane) are already in the layout the GEMM
-// reads (same trans4_ns convert that feeds gemm_moe_q5_k_f32_ns), so only the scale
-// is rebuilt here.
-//
-// One work-item per (row, superblock, expert); each emits 8 sub-blocks.
-// ---------------------------------------------------------------------------
-kernel void kernel_moe_expand_scale_q5_K(
-    global const uchar * src_s,     // [expert][row][superblock][12]
-    global const half  * src_d,     // [expert][superblock][row]
-    global const half  * src_dm,    // [expert][superblock][row]
-    global       half  * dst_scale, // [expert][row][32block][2]
-    global       half  * dst_min,   // [expert][row][32block]
-    int ne00,
-    int ne01
-) {
-    int row = get_global_id(0);
-    int sb  = get_global_id(1);   // superblock index along K
-    int e   = get_global_id(2);
-    if (row >= ne01) { return; }
-
-    long nsb    = ne00 / 256;     // superblocks per row
-    long nblk32 = ne00 / 32;      // 32-blocks per row
-
-    float d  = (float)src_d [((long)e*nsb + sb)*ne01 + row];
-    float dm = (float)src_dm[((long)e*nsb + sb)*ne01 + row];
-
-    __global const uchar * sc = src_s + ((long)e*ne01 + row)*nsb*12 + (long)sb*12;
-
-    for (int j = 0; j < 8; ++j) {
-        uchar sv, mn;
-        // get_scale_min_k4 (6-bit packed scale/min for sub-block j of 8)
-        if (j < 4) {
-            sv = sc[j]   & 63;
-            mn = sc[j+4] & 63;
-        } else {
-            sv = (sc[j+4] & 0x0F) | ((sc[j-4] & 0xC0) >> 2);
-            mn = ((sc[j+4] >> 4) & 0x0F) | ((sc[j]   & 0xC0) >> 2);
-        }
-        long sub   = (long)sb*8 + j;
-        long sbase = (((long)e*ne01 + row)*nblk32 + sub) * 2;
-        half s_val = (half)(d  * (float)sv);
-        dst_scale[sbase + 0] = s_val;
-        dst_scale[sbase + 1] = s_val;
-        dst_min[((long)e*ne01 + row)*nblk32 + sub] = (half)(dm * (float)mn);
-    }
-}

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

@@ -1,186 +0,0 @@
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-#pragma OPENCL EXTENSION cl_khr_subgroups : enable
-#ifdef cl_khr_integer_dot_product
-#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
-#endif
-
-#define TILESIZE_M 64
-#define TILESIZE_N 32
-
-// 2*mxfp4_value as signed int8, packed 4 codes per uint. Divergent nibble
-// lookups read a __constant *uint* array + shift, never a byte array
-// (byte-indexed __constant loads serialize on Adreno and are far slower).
-//   idx 0-3:   0,  1,  2,  3   = 0x03020100
-//   idx 4-7:   4,  6,  8, 12   = 0x0C080604
-//   idx 8-11:  0, -1, -2, -3   = 0xFDFEFF00   (-1=0xFF,-2=0xFE,-3=0xFD)
-//   idx 12-15:-4, -6, -8,-12   = 0xF4F8FAFC   (-4=0xFC,-6=0xFA,-8=0xF8,-12=0xF4)
-__constant uint mxfp4_i8x4[4] = {
-    0x03020100u, 0x0C080604u, 0xFDFEFF00u, 0xF4F8FAFCu
-};
-inline uint mxfp4_code(uint n) {
-    return (mxfp4_i8x4[n >> 2] >> ((n & 3u) * 8u)) & 0xFFu;
-}
-// 4 nibbles in the low 16 bits of u -> 4 codebook int8, packed for dp4a.
-inline uint mxfp4_pack(ushort u) {
-    return  mxfp4_code((uint)( u        & 0xF))
-         | (mxfp4_code((uint)((u >>  4) & 0xF)) <<  8)
-         | (mxfp4_code((uint)((u >>  8) & 0xF)) << 16)
-         | (mxfp4_code((uint)((u >> 12) & 0xF)) << 24);
-}
-
-static inline float e8m0_to_fp32(uchar x) {
-    int bits;
-    bits = (x == 0) ? 0x00400000 : ((uint) x << 23);
-    return as_float(bits);
-}
-
-// One token's dp4a dot (8 uints = 32 K elems) + mxfp4 block-scale epilogue.
-// blk_scale already carries the 0.5 factor (== 0.5 * 2^e).
-#define MOE_MXFP4_DP4A_T(t) do {                                     \
-        int raw = 0;                                                 \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[0], sh_qa[t][0], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[1], sh_qa[t][1], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[2], sh_qa[t][2], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[3], sh_qa[t][3], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[4], sh_qa[t][4], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[5], sh_qa[t][5], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[6], sh_qa[t][6], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[7], sh_qa[t][7], raw); \
-        acc[t] += blk_scale * (float)sh_d[t] * (float)raw;           \
-    } while (0)
-
-__attribute__((qcom_wave_pair_mode(1)))
-kernel void kernel_gemm_moe_mxfp4_q8_1_dp4a(
-        __read_only  image1d_buffer_t src0_q,    // mxfp4 codes (transposed, packed nibbles)
-        __global     uchar *          src0_e,    // e8m0 per-32-block scale
-        __global     uint *           src1_qa,   // q8_1 activations: int8 quants (as uint, 4/elem)
-        __global     half *           src1_da,   // q8_1 per-block scale  [tok_slot * ne00/32]
-        __global     uint *           src2,      // post-router (orig out positions)
-        __global     ushort *         src2_emap, // tile -> expert id
-        __write_only image1d_buffer_t dst,
-        __global     int *            total_tiles,
-        uint ne00,
-        uint ne01,
-        int  is_ragged                           // 1: compute only real tokens per tile
-) {
-    const uint block_id_m = get_global_id(1); // m_tile
-    const uint block_id_n = get_global_id(2); // n_tile
-
-    if (block_id_n >= total_tiles[0]) {
-        return;
-    }
-
-    const uint lid = get_local_id(0);          // 0..63, == this WI's output row in the M-tile
-
-    const ushort expert_id = src2_emap[block_id_n];
-    const uint   row = block_id_m * TILESIZE_M;
-    const uint   col = block_id_n * TILESIZE_N;
-
-    const uint num_blocks = ne00 >> 5;          // blocks-of-32 per token
-    const uint row_idx    = row + lid;
-
-    const uint ne00_u = ne00 >> 2;   // ne00 in uint (int8x4) units
-
-    __local uint sh_qa[TILESIZE_N][8]; // 32 tokens x 8 uints (32 int8) = 1 KiB
-    __local half sh_d[TILESIZE_N];
-
-    // Real token count for this tile.
-    // Real tokens are packed contiguously at the tile start; padded slots hold
-    // 0xFFFFFFFF (only the last tile of each expert is partial). is_ragged skips
-    // the dp4a/staging/scatter for padded slots; is_ragged==0 forces n_real=32.
-    __local uint sh_src2[TILESIZE_N];
-    __local int  sh_nreal;
-    if (lid < TILESIZE_N) {
-        sh_src2[lid] = src2[col + lid];
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-    if (lid == 0) {
-        int nr = TILESIZE_N;
-        if (is_ragged) {
-            nr = 0;
-            #pragma unroll
-            for (int t = 0; t < TILESIZE_N; ++t) {
-                if (sh_src2[t] != 0xFFFFFFFFu) ++nr;
-            }
-        }
-        sh_nreal = nr;
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-    const int n_real = sh_nreal;
-
-    float acc[TILESIZE_N];
-    #pragma unroll
-    for (int t = 0; t < TILESIZE_N; ++t) acc[t] = 0.0f;
-
-    for (uint step = 0; step < ne00; step += 32) {
-        const uint sub = step >> 5;        // 32-block index along K
-
-        // e8m0 block scale for this WI's row, this 32-block (folded x0.5)
-        const uint e_offset = row_idx + sub * ne01 + expert_id * num_blocks * ne01;
-        const float blk_scale = 0.5f * e8m0_to_fp32(src0_e[e_offset]);
-
-        // repack this WI's 32 weight nibbles into 8 dp4a uints
-        const uint qoff0 = row + ((ne01 * step) >> 3)        + ((expert_id * ne00 * ne01) >> 3);
-        const uint qoff1 = row + ((ne01 * (step + 16)) >> 3) + ((expert_id * ne00 * ne01) >> 3);
-        const uint r0 = read_imageui(src0_q, qoff0 + lid).x;
-        const uint r1 = read_imageui(src0_q, qoff0 + lid + ne01).x;
-        const uint r2 = read_imageui(src0_q, qoff1 + lid).x;
-        const uint r3 = read_imageui(src0_q, qoff1 + lid + ne01).x;
-        uint qw[8];
-        qw[0] = mxfp4_pack((ushort)(r0));        qw[1] = mxfp4_pack((ushort)(r0 >> 16));
-        qw[2] = mxfp4_pack((ushort)(r1));        qw[3] = mxfp4_pack((ushort)(r1 >> 16));
-        qw[4] = mxfp4_pack((ushort)(r2));        qw[5] = mxfp4_pack((ushort)(r2 >> 16));
-        qw[6] = mxfp4_pack((ushort)(r3));        qw[7] = mxfp4_pack((ushort)(r3 >> 16));
-
-        // cooperatively stage the n_real-token x 32-K int8 activations
-        const uint stage_lim = (uint)n_real * 8;
-        for (uint idx = lid; idx < stage_lim; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u];
-        }
-        if (lid < (uint)n_real) {
-            sh_d[lid] = src1_da[(col + lid) * num_blocks + sub];
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-
-        // Full tiles keep the fully-unrolled 32-wide loop; partial tiles run only n_real
-        if (n_real == TILESIZE_N) {
-            #pragma unroll
-            for (int t = 0; t < TILESIZE_N; ++t) { MOE_MXFP4_DP4A_T(t); }
-        } else {
-            #pragma unroll 4
-            for (int t = 0; t < n_real; ++t) { MOE_MXFP4_DP4A_T(t); }
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-    }
-
-    if (row_idx >= ne01) {
-        return;
-    }
-
-    // scatter results to original output rows (reuse sh_src2 from the top)
-    __local uint out_idx[TILESIZE_N];
-    if (lid < TILESIZE_N) {
-        uint idx = sh_src2[lid];
-        if (idx == 0xFFFFFFFF) {
-            idx = sh_src2[0];
-        }
-        out_idx[lid] = idx * ne01;
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-
-    const uint m_offset = row + lid;
-    if (n_real == TILESIZE_N) {
-        #pragma unroll
-        for (int t = 1; t < TILESIZE_N; ++t) {
-            write_imagef(dst, out_idx[t] + m_offset, acc[t]);
-        }
-        barrier(CLK_GLOBAL_MEM_FENCE);
-        write_imagef(dst, out_idx[0] + m_offset, acc[0]);
-    } else {
-        for (int t = 0; t < n_real; ++t) {
-            write_imagef(dst, out_idx[t] + m_offset, acc[t]);
-        }
-    }
-}

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

@@ -1,165 +0,0 @@
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-#pragma OPENCL EXTENSION cl_khr_subgroups : enable
-#ifdef cl_khr_integer_dot_product
-#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
-#endif
-
-#define TILESIZE_M 64
-#define TILESIZE_N 32
-
-// Expand the 4 nibbles held in the low 16 bits of `u` into 4 bytes (one nibble
-// per byte, value 0..15), packed for the int8 dp4a. The -8 zero-point is applied
-// in the epilogue via the activation sum term (cheaper than biasing every byte).
-#define EXP4(u)  ( ((uint)((u) & 0x000Fu))        | \
-                  (((uint)((u) & 0x00F0u)) << 4)  | \
-                  (((uint)((u) & 0x0F00u)) << 8)  | \
-                  (((uint)((u) & 0xF000u)) << 12) )
-
-// One token's dp4a dot (8 uints = 32 K elems) + q4_0 scale/zero-point epilogue.
-#define MOE_Q40_DP4A_T(t) do {                                       \
-        int raw = 0;                                                 \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[0], sh_qa[t][0], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[1], sh_qa[t][1], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[2], sh_qa[t][2], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[3], sh_qa[t][3], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[4], sh_qa[t][4], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[5], sh_qa[t][5], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[6], sh_qa[t][6], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[7], sh_qa[t][7], raw); \
-        acc[t] += d_val * ((float)sh_d[t] * (float)raw - 8.0f * (float)sh_s[t]); \
-    } while (0)
-
-__attribute__((qcom_wave_pair_mode(1)))
-kernel void kernel_gemm_moe_q4_0_q8_1_dp4a(
-        __read_only  image1d_buffer_t src0_q,   // q4_0 weights (transposed, packed nibbles)
-        __global     half *           src0_d,   // per-32-block scale
-        __global     uint *           src1_qa,  // q8_1 activations: int8 quants (as uint, 4/elem)
-        __global     half *           src1_da,  // q8_1 per-block scale  [tok_slot * ne00/32]
-        __global     half *           src1_sa,  // q8_1 per-block sum*d  [tok_slot * ne00/32]
-        __global     uint *           src2,     // post-router (orig out positions)
-        __global     ushort *         src2_emap,// tile -> expert id
-        __write_only image1d_buffer_t dst,
-        __global     int *            total_tiles,
-        uint ne00,
-        uint ne01,
-        int  is_ragged                          // 1: compute only real tokens per tile
-) {
-    const uint block_id_m = get_global_id(1); // m_tile
-    const uint block_id_n = get_global_id(2); // n_tile
-
-    if (block_id_n >= total_tiles[0]) {
-        return;
-    }
-
-    const uint lid = get_local_id(0);          // 0..63, == this WI's output row in the M-tile
-
-    const ushort expert_id = src2_emap[block_id_n];
-    const uint   row = block_id_m * TILESIZE_M;
-    const uint   col = block_id_n * TILESIZE_N;
-
-    const uint num_blocks = ne00 >> 5;          // blocks-of-32 per token
-    const uint row_idx    = row + lid;
-
-    const uint ne00_u = ne00 >> 2;   // ne00 in uint (int8x4) units
-
-    __local uint sh_qa[TILESIZE_N][8]; // 32 tokens x 8 uints (32 int8) = 1 KiB
-    __local half sh_d[TILESIZE_N];
-    __local half sh_s[TILESIZE_N];
-
-    // Real-token count for this tile
-    __local uint sh_src2[TILESIZE_N];
-    __local int  sh_nreal;
-    if (lid < TILESIZE_N) {
-        sh_src2[lid] = src2[col + lid];
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-    if (lid == 0) {
-        int nr = TILESIZE_N;
-        if (is_ragged) {
-            nr = 0;
-            #pragma unroll
-            for (int t = 0; t < TILESIZE_N; ++t) {
-                if (sh_src2[t] != 0xFFFFFFFFu) ++nr;
-            }
-        }
-        sh_nreal = nr;
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-    const int n_real = sh_nreal;
-
-    float acc[TILESIZE_N];
-    #pragma unroll
-    for (int t = 0; t < TILESIZE_N; ++t) acc[t] = 0.0f;
-
-    for (uint step = 0; step < ne00; step += 32) {
-        const uint sub = step >> 5;        // 32-block index along K
-
-        // per-32-block scale for this WI's row
-        const uint d_offset = row_idx + sub * ne01 + expert_id * num_blocks * ne01;
-        const float d_val = (float)src0_d[d_offset];
-
-        // repack this WI's 32 weight nibbles into 8 dp4a uints
-        const uint qoff0 = row + ((ne01 * step) >> 3)        + ((expert_id * ne00 * ne01) >> 3);
-        const uint qoff1 = row + ((ne01 * (step + 16)) >> 3) + ((expert_id * ne00 * ne01) >> 3);
-        const uint r0 = read_imageui(src0_q, qoff0 + lid).x;
-        const uint r1 = read_imageui(src0_q, qoff0 + lid + ne01).x;
-        const uint r2 = read_imageui(src0_q, qoff1 + lid).x;
-        const uint r3 = read_imageui(src0_q, qoff1 + lid + ne01).x;
-        uint qw[8];
-        qw[0] = EXP4(r0);        qw[1] = EXP4(r0 >> 16);
-        qw[2] = EXP4(r1);        qw[3] = EXP4(r1 >> 16);
-        qw[4] = EXP4(r2);        qw[5] = EXP4(r2 >> 16);
-        qw[6] = EXP4(r3);        qw[7] = EXP4(r3 >> 16);
-
-        // cooperatively stage the n_real-token x 32-K int8 activations
-        const uint stage_lim = (uint)n_real * 8;
-        for (uint idx = lid; idx < stage_lim; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u];
-        }
-        if (lid < (uint)n_real) {
-            sh_d[lid] = src1_da[(col + lid) * num_blocks + sub];
-            sh_s[lid] = src1_sa[(col + lid) * num_blocks + sub];
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-
-        if (n_real == TILESIZE_N) {
-            #pragma unroll
-            for (int t = 0; t < TILESIZE_N; ++t) { MOE_Q40_DP4A_T(t); }
-        } else {
-            #pragma unroll 4
-            for (int t = 0; t < n_real; ++t) { MOE_Q40_DP4A_T(t); }
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-    }
-
-    if (row_idx >= ne01) {
-        return;
-    }
-
-    // scatter results to original output rows (reuse sh_src2 from the top)
-    __local uint out_idx[TILESIZE_N];
-    if (lid < TILESIZE_N) {
-        uint idx = sh_src2[lid];
-        if (idx == 0xFFFFFFFF) {
-            idx = sh_src2[0];
-        }
-        out_idx[lid] = idx * ne01;
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-
-    const uint m_offset = row + lid;
-    if (n_real == TILESIZE_N) {
-        #pragma unroll
-        for (int t = 1; t < TILESIZE_N; ++t) {
-            write_imagef(dst, out_idx[t] + m_offset, acc[t]);
-        }
-        barrier(CLK_GLOBAL_MEM_FENCE);
-        write_imagef(dst, out_idx[0] + m_offset, acc[0]);
-    } else {
-        for (int t = 0; t < n_real; ++t) {
-            write_imagef(dst, out_idx[t] + m_offset, acc[t]);
-        }
-    }
-}

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

@@ -1,202 +0,0 @@
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-#pragma OPENCL EXTENSION cl_khr_subgroups : enable
-#ifdef cl_khr_integer_dot_product
-#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
-#endif
-
-// q4_K subblock (32 elems): w_i = scale*q_i - minv, q_i in [0,15], scale =
-// d_super*sv6, minv = dmin_super*mn6. With activation block (a_d, a_s, qa[32]):
-//   Sum_i w_i * a_i = scale * a_d * dp4a(q, qa) - minv * a_s
-// where a_s = a_d * Sum(qa) (the q8_1 "s" field)
-
-#define TILESIZE_M 64
-#define TILESIZE_N 32
-#define QK_K 256
-#define K_SCALE_SIZE 12
-
-inline void get_scale_min_k4(
-    int j,
-    global const uchar * q,
-    uchar * d,
-    uchar * m
-) {
-    if (j < 4) {
-        *d = q[j]   & 63;
-        *m = q[j+4] & 63;
-    } else {
-        *d = (q[j+4] & 0x0F) | ((q[j-4] & 0xC0) >> 2);
-        *m = ((q[j+4] >> 4) & 0x0F) | ((q[j]   & 0xC0) >> 2);
-    }
-}
-
-// Expand the 4 nibbles held in the low 16 bits of `u` into 4 bytes (one nibble
-// per byte, value 0..15), packed for the int8 dp4a.
-#define EXP4(u)  ( ((uint)((u) & 0x000Fu))        | \
-                  (((uint)((u) & 0x00F0u)) << 4)  | \
-                  (((uint)((u) & 0x0F00u)) << 8)  | \
-                  (((uint)((u) & 0xF000u)) << 12) )
-
-// One token's dp4a dot (8 uints = 32 K elems) + q4_K scale/min epilogue into acc[t].
-#define MOE_Q4K_DP4A_T(t) do {                                       \
-        int raw = 0;                                                 \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[0], sh_qa[t][0], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[1], sh_qa[t][1], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[2], sh_qa[t][2], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[3], sh_qa[t][3], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[4], sh_qa[t][4], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[5], sh_qa[t][5], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[6], sh_qa[t][6], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[7], sh_qa[t][7], raw); \
-        acc[t] += scale * (float)sh_d[t] * (float)raw - minv * (float)sh_s[t]; \
-    } while (0)
-
-__attribute__((qcom_wave_pair_mode(1)))
-kernel void kernel_gemm_moe_q4_k_q8_1_dp4a(
-        __read_only  image1d_buffer_t src0_q,   // q4_K weights (transposed, packed nibbles)
-        __global     half *           src0_d,   // per-superblock scale
-        __global     half *           src0_dm,  // per-superblock min
-        __global     uchar *          src0_s,   // 6-bit scale/min codes
-        __global     uint *           src1_qa,  // q8_1 activations: int8 quants (as uint, 4/elem)
-        __global     half *           src1_da,  // q8_1 per-block scale  [tok_slot * ne00/32]
-        __global     half *           src1_sa,  // q8_1 per-block sum*d  [tok_slot * ne00/32]
-        __global     uint *           src2,     // post-router (orig out positions)
-        __global     ushort *         src2_emap,// tile -> expert id
-        __write_only image1d_buffer_t dst,
-        __global     int *            total_tiles,
-        uint ne00,
-        uint ne01,
-        int  is_ragged                          // 1: compute only real tokens per tile
-) {
-    const uint block_id_m = get_global_id(1); // m_tile
-    const uint block_id_n = get_global_id(2); // n_tile
-
-    if (block_id_n >= total_tiles[0]) {
-        return;
-    }
-
-    const uint lid = get_local_id(0);          // 0..63, == this WI's output row in the M-tile
-
-    const ushort expert_id = src2_emap[block_id_n];
-    const uint   row = block_id_m * TILESIZE_M;
-    const uint   col = block_id_n * TILESIZE_N;
-
-    const uint num_superblocks = ne00 / QK_K;
-    const uint scales_per_row  = num_superblocks * K_SCALE_SIZE;
-    const uint row_idx         = row + lid;
-
-    const uint ne00_u  = ne00 >> 2;   // ne00 in uint (int8x4) units
-    const uint ne00_b  = ne00 >> 5;   // blocks-of-32 per token
-
-    __local uint sh_qa[TILESIZE_N][8]; // 32 tokens x 8 uints (32 int8) = 1 KiB
-    __local half sh_d[TILESIZE_N];
-    __local half sh_s[TILESIZE_N];
-
-    // Real token count for this tile
-    __local uint sh_src2[TILESIZE_N];
-    __local int  sh_nreal;
-    if (lid < TILESIZE_N) {
-        sh_src2[lid] = src2[col + lid];
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-    if (lid == 0) {
-        int nr = TILESIZE_N;
-        if (is_ragged) {
-            nr = 0;
-            #pragma unroll
-            for (int t = 0; t < TILESIZE_N; ++t) {
-                if (sh_src2[t] != 0xFFFFFFFFu) ++nr;
-            }
-        }
-        sh_nreal = nr;
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-    const int n_real = sh_nreal;
-
-    float acc[TILESIZE_N];
-    #pragma unroll
-    for (int t = 0; t < TILESIZE_N; ++t) acc[t] = 0.0f;
-
-    for (uint step = 0; step < ne00; step += 32) {
-        const uint sub = step >> 5;        // subblock index along K
-        const uint sb  = sub >> 3;         // superblock index
-        const uint j   = sub & 7;          // subblock within superblock
-
-        // --- weight scale / min for this WI's row, this subblock ---
-        const uint d_offset = row + sb * ne01 + expert_id * num_superblocks * ne01 + lid;
-        const float d_val  = (float)src0_d[d_offset];
-        const float dm_val = (float)src0_dm[d_offset];
-
-        global const uchar * sc = src0_s + (expert_id * ne01 + row_idx) * scales_per_row + sb * K_SCALE_SIZE;
-        uchar sv, mn;
-        get_scale_min_k4(j, sc, &sv, &mn);
-        const float scale = d_val  * (float)sv;
-        const float minv  = dm_val * (float)mn;
-
-        // --- repack this WI's 32 weight nibbles into 8 dp4a uints ---
-        const uint qoff0 = row + ((ne01 * step) >> 3)        + ((expert_id * ne00 * ne01) >> 3);
-        const uint qoff1 = row + ((ne01 * (step + 16)) >> 3) + ((expert_id * ne00 * ne01) >> 3);
-        const uint r0 = read_imageui(src0_q, qoff0 + lid).x;
-        const uint r1 = read_imageui(src0_q, qoff0 + lid + ne01).x;
-        const uint r2 = read_imageui(src0_q, qoff1 + lid).x;
-        const uint r3 = read_imageui(src0_q, qoff1 + lid + ne01).x;
-        uint qw[8];
-        qw[0] = EXP4(r0);        qw[1] = EXP4(r0 >> 16);
-        qw[2] = EXP4(r1);        qw[3] = EXP4(r1 >> 16);
-        qw[4] = EXP4(r2);        qw[5] = EXP4(r2 >> 16);
-        qw[6] = EXP4(r3);        qw[7] = EXP4(r3 >> 16);
-
-        // --- cooperatively stage the n_real-token x 32-K int8 activations to LDS ---
-        const uint stage_lim = (uint)n_real * 8;
-        for (uint idx = lid; idx < stage_lim; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u];
-        }
-        if (lid < (uint)n_real) {
-            sh_d[lid] = src1_da[(col + lid) * ne00_b + sub];
-            sh_s[lid] = src1_sa[(col + lid) * ne00_b + sub];
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-
-        // dp4a - each real token sum over 8 uints (32 K), then scale/min
-        // Full tiles keep the fully-unrolled 32-wide loop;
-        // partial tiles run only n_real (saves the padded-slot dp4a + staging).
-        if (n_real == TILESIZE_N) {
-            #pragma unroll
-            for (int t = 0; t < TILESIZE_N; ++t) { MOE_Q4K_DP4A_T(t); }
-        } else {
-            #pragma unroll 4
-            for (int t = 0; t < n_real; ++t) { MOE_Q4K_DP4A_T(t); }
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-    }
-
-    if (row_idx >= ne01) {
-        return;
-    }
-
-    // scatter results to original output rows
-    __local uint out_idx[TILESIZE_N];
-    if (lid < TILESIZE_N) {
-        uint idx = sh_src2[lid];
-        if (idx == 0xFFFFFFFF) {
-            idx = sh_src2[0];
-        }
-        out_idx[lid] = idx * ne01;
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-
-    const uint m_offset = row + lid;
-    if (n_real == TILESIZE_N) {
-        #pragma unroll
-        for (int t = 1; t < TILESIZE_N; ++t) {
-            write_imagef(dst, out_idx[t] + m_offset, acc[t]);
-        }
-        barrier(CLK_GLOBAL_MEM_FENCE);
-        write_imagef(dst, out_idx[0] + m_offset, acc[0]);
-    } else {
-        for (int t = 0; t < n_real; ++t) {
-            write_imagef(dst, out_idx[t] + m_offset, acc[t]);
-        }
-    }
-}

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

@@ -1,196 +0,0 @@
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-#pragma OPENCL EXTENSION cl_khr_subgroups : enable
-#ifdef cl_khr_integer_dot_product
-#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
-#endif
-
-#define TILESIZE_N 32
-#define QK_K 256
-
-// 4 nibbles in the low 16 bits of `u` -> 4 bytes (value 0..15, in bits 0-3).
-#define EXP4(u)  ( ((uint)((u) & 0x000Fu))        | \
-                  (((uint)((u) & 0x00F0u)) << 4)  | \
-                  (((uint)((u) & 0x0F00u)) << 8)  | \
-                  (((uint)((u) & 0xF000u)) << 12) )
-
-// 4 2-bit highs in byte `b` (8 bits) -> 4 bytes, value 0..3 in bits 4-5
-// (pre-multiplied by 16 so it ORs with the EXP4 nibble to form q6 in 0..63).
-#define EXP2(b)  ( (((uint)((b) & 0x03u)) << 4)   | \
-                  (((uint)((b) & 0x0Cu)) << 10)  | \
-                  (((uint)((b) & 0x30u)) << 16)  | \
-                  (((uint)((b) & 0xC0u)) << 22) )
-
-// q6 (0..63, bits 0-5 of each byte) -> (q6-32) as a signed int8 per byte.
-// Flipping bit5 subtracts 32 in 6-bit two's complement; then replicate bit5
-// into bits 6-7 to sign-extend to int8. Per-byte, no inter-byte carry.
-inline uint SIGN6(uint q6p) {
-    uint x = q6p ^ 0x20202020u;
-    uint s = x & 0x20202020u;
-    return x | (s << 1) | (s << 2);
-}
-
-inline int dp4a_q6(uint qw0, uint qw1, uint qw2, uint qw3,
-                   uint a0, uint a1, uint a2, uint a3) {
-    int raw = 0;
-    raw = dot_acc_sat_4x8packed_ss_int(qw0, a0, raw);
-    raw = dot_acc_sat_4x8packed_ss_int(qw1, a1, raw);
-    raw = dot_acc_sat_4x8packed_ss_int(qw2, a2, raw);
-    raw = dot_acc_sat_4x8packed_ss_int(qw3, a3, raw);
-    return raw;
-}
-
-// One token's q6_K dp4a dot (two halves, per-16 scales) + epilogue into acc[t].
-#define MOE_Q6K_DP4A_T(t) do {                                                                            \
-        const int raw1 = dp4a_q6(qw[0], qw[1], qw[2], qw[3], sh_qa[t][0], sh_qa[t][1], sh_qa[t][2], sh_qa[t][3]); \
-        const int raw2 = dp4a_q6(qw[4], qw[5], qw[6], qw[7], sh_qa[t][4], sh_qa[t][5], sh_qa[t][6], sh_qa[t][7]); \
-        const float a_d = (float)sh_d[t];                                                                 \
-        acc[t] += scale0 * a_d * (float)raw1 + scale1 * a_d * (float)raw2;                                \
-    } while (0)
-
-__attribute__((qcom_wave_pair_mode(1)))
-kernel void kernel_gemm_moe_q6_k_q8_1_dp4a(
-        __read_only  image1d_buffer_t src0_ql,   // q6_K low nibbles (image, q4_K-style layout)
-        __global     uint *           src0_qh,   // q6_K high 2-bit (16 elems/uint)
-        __global     char *           src0_s,    // int8 scales (one per 16 elems)
-        __global     half *           src0_d,    // per-superblock scale
-        __global     uint *           src1_qa,   // q8_1 activations int8 (as uint, 4/elem)
-        __global     half *           src1_da,   // q8_1 per-block scale [tok_slot * ne00/32]
-        __global     uint *           src2,      // post-router (orig out positions)
-        __global     ushort *         src2_emap, // tile -> expert id
-        __write_only image1d_buffer_t dst,
-        __global     int *            total_tiles,
-        uint ne00,
-        uint ne01,
-        int  is_ragged                         // 1: compute only real tokens per tile
-) {
-    const uint block_id_m = get_global_id(1);
-    const uint block_id_n = get_global_id(2);
-
-    if (block_id_n >= total_tiles[0]) {
-        return;
-    }
-
-    const uint lid = get_local_id(0);          // 0..63 -> row within M-tile
-
-    const ushort expert_id = src2_emap[block_id_n];
-    const uint   row = block_id_m * 64;
-    const uint   col = block_id_n * TILESIZE_N;
-
-    const uint num_superblocks = ne00 / QK_K;
-    const uint scales_per_row  = num_superblocks * 16;
-    const uint row_idx         = row + lid;
-
-    const uint ne00_u = ne00 >> 2;
-    const uint ne00_b = ne00 >> 5;
-
-    __local uint sh_qa[TILESIZE_N][8];
-    __local half sh_d[TILESIZE_N];
-
-    // Real token count for this tile
-    __local uint sh_src2[TILESIZE_N];
-    __local int  sh_nreal;
-    if (lid < TILESIZE_N) {
-        sh_src2[lid] = src2[col + lid];
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-    if (lid == 0) {
-        int nr = TILESIZE_N;
-        if (is_ragged) {
-            nr = 0;
-            #pragma unroll
-            for (int t = 0; t < TILESIZE_N; ++t) {
-                if (sh_src2[t] != 0xFFFFFFFFu) ++nr;
-            }
-        }
-        sh_nreal = nr;
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-    const int n_real = sh_nreal;
-
-    float acc[TILESIZE_N];
-    #pragma unroll
-    for (int t = 0; t < TILESIZE_N; ++t) acc[t] = 0.0f;
-
-    for (uint step = 0; step < ne00; step += 32) {
-        const uint sub = step >> 5;
-        const uint sb  = sub >> 3;
-        const uint j   = sub & 7;
-
-        const float d_val = (float)src0_d[row + sb * ne01 + expert_id * num_superblocks * ne01 + lid];
-        global const char * sc = src0_s + (expert_id * ne01 + row_idx) * scales_per_row + sb * 16;
-        const float scale0 = d_val * (float)sc[j * 2];
-        const float scale1 = d_val * (float)sc[j * 2 + 1];
-
-        // high bits: one uint covers 16 elems; first/second 16 of this 32-block
-        const uint qh_base = row + (sub * 2) * ne01 + expert_id * (num_superblocks * 16) * ne01 + lid;
-        const uint qh1 = src0_qh[qh_base];
-        const uint qh2 = src0_qh[qh_base + ne01];
-
-        // low nibbles: same image layout as q4_K (8 ushorts over the 32 K)
-        const uint qoff0 = row + ((ne01 * step) >> 3)        + ((expert_id * ne00 * ne01) >> 3);
-        const uint qoff1 = row + ((ne01 * (step + 16)) >> 3) + ((expert_id * ne00 * ne01) >> 3);
-        const uint r0 = read_imageui(src0_ql, qoff0 + lid).x;
-        const uint r1 = read_imageui(src0_ql, qoff0 + lid + ne01).x;
-        const uint r2 = read_imageui(src0_ql, qoff1 + lid).x;
-        const uint r3 = read_imageui(src0_ql, qoff1 + lid + ne01).x;
-
-        uint qw[8];
-        qw[0] = SIGN6(EXP4(r0)       | EXP2((qh1)       & 0xFFu));
-        qw[1] = SIGN6(EXP4(r0 >> 16) | EXP2((qh1 >> 8)  & 0xFFu));
-        qw[2] = SIGN6(EXP4(r1)       | EXP2((qh1 >> 16) & 0xFFu));
-        qw[3] = SIGN6(EXP4(r1 >> 16) | EXP2((qh1 >> 24) & 0xFFu));
-        qw[4] = SIGN6(EXP4(r2)       | EXP2((qh2)       & 0xFFu));
-        qw[5] = SIGN6(EXP4(r2 >> 16) | EXP2((qh2 >> 8)  & 0xFFu));
-        qw[6] = SIGN6(EXP4(r3)       | EXP2((qh2 >> 16) & 0xFFu));
-        qw[7] = SIGN6(EXP4(r3 >> 16) | EXP2((qh2 >> 24) & 0xFFu));
-
-        const uint stage_lim = (uint)n_real * 8;
-        for (uint idx = lid; idx < stage_lim; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u];
-        }
-        if (lid < (uint)n_real) {
-            sh_d[lid] = src1_da[(col + lid) * ne00_b + sub];
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-
-        // Full tiles keep the fully-unrolled 32-wide loop; partial tiles run n_real.
-        if (n_real == TILESIZE_N) {
-            #pragma unroll
-            for (int t = 0; t < TILESIZE_N; ++t) { MOE_Q6K_DP4A_T(t); }
-        } else {
-            #pragma unroll 4
-            for (int t = 0; t < n_real; ++t) { MOE_Q6K_DP4A_T(t); }
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-    }
-
-    if (row_idx >= ne01) {
-        return;
-    }
-
-    __local uint out_idx[TILESIZE_N];
-    if (lid < TILESIZE_N) {
-        uint idx = sh_src2[lid];
-        if (idx == 0xFFFFFFFF) {
-            idx = sh_src2[0];
-        }
-        out_idx[lid] = idx * ne01;
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-
-    const uint m_offset = row + lid;
-    if (n_real == TILESIZE_N) {
-        #pragma unroll
-        for (int t = 1; t < TILESIZE_N; ++t) {
-            write_imagef(dst, out_idx[t] + m_offset, acc[t]);
-        }
-        barrier(CLK_GLOBAL_MEM_FENCE);
-        write_imagef(dst, out_idx[0] + m_offset, acc[0]);
-    } else {
-        for (int t = 0; t < n_real; ++t) {
-            write_imagef(dst, out_idx[t] + m_offset, acc[t]);
-        }
-    }
-}

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

@@ -1,221 +0,0 @@
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-#pragma OPENCL EXTENSION cl_khr_subgroups : enable
-#pragma OPENCL EXTENSION cl_qcom_subgroup_uniform_load: enable
-#pragma OPENCL EXTENSION cl_qcom_subgroup_constant_load: enable
-#pragma OPENCL EXTENSION cl_qcom_extra_vector_types : enable
-
-#define TILESIZE_K 16
-#define TILESIZE_M 64
-#define TILESIZE_N 32
-
-// q8_0: 16 signed int8 weights (one uint4 = 16 chars) -> half16, scaled.
-#define dequantize_q8_0(q4, a_f16, scale) \
-    a_f16 = convert_half16(as_char16(q4)) * scale;
-
-#define dotx16_reduce8(a_reg, b_lm, c_reg, lm_offset) \
-    acc.s0 = dot(a_reg.s0123, b_lm[lm_offset + 0]); \
-    acc.s1 = dot(a_reg.s0123, b_lm[lm_offset + 1]); \
-    acc.s2 = dot(a_reg.s0123, b_lm[lm_offset + 2]); \
-    acc.s3 = dot(a_reg.s0123, b_lm[lm_offset + 3]); \
-    acc.s4 = dot(a_reg.s0123, b_lm[lm_offset + 4]); \
-    acc.s5 = dot(a_reg.s0123, b_lm[lm_offset + 5]); \
-    acc.s6 = dot(a_reg.s0123, b_lm[lm_offset + 6]); \
-    acc.s7 = dot(a_reg.s0123, b_lm[lm_offset + 7]); \
-    acc.s8 = dot(a_reg.s0123, b_lm[lm_offset + 8]); \
-    acc.s9 = dot(a_reg.s0123, b_lm[lm_offset + 9]); \
-    acc.sa = dot(a_reg.s0123, b_lm[lm_offset + 10]); \
-    acc.sb = dot(a_reg.s0123, b_lm[lm_offset + 11]); \
-    acc.sc = dot(a_reg.s0123, b_lm[lm_offset + 12]); \
-    acc.sd = dot(a_reg.s0123, b_lm[lm_offset + 13]); \
-    acc.se = dot(a_reg.s0123, b_lm[lm_offset + 14]); \
-    acc.sf = dot(a_reg.s0123, b_lm[lm_offset + 15]); \
-    acc.s0 += dot(a_reg.s4567, b_lm[lm_offset + 32]); \
-    acc.s1 += dot(a_reg.s4567, b_lm[lm_offset + 33]); \
-    acc.s2 += dot(a_reg.s4567, b_lm[lm_offset + 34]); \
-    acc.s3 += dot(a_reg.s4567, b_lm[lm_offset + 35]); \
-    acc.s4 += dot(a_reg.s4567, b_lm[lm_offset + 36]); \
-    acc.s5 += dot(a_reg.s4567, b_lm[lm_offset + 37]); \
-    acc.s6 += dot(a_reg.s4567, b_lm[lm_offset + 38]); \
-    acc.s7 += dot(a_reg.s4567, b_lm[lm_offset + 39]); \
-    acc.s8 += dot(a_reg.s4567, b_lm[lm_offset + 40]); \
-    acc.s9 += dot(a_reg.s4567, b_lm[lm_offset + 41]); \
-    acc.sa += dot(a_reg.s4567, b_lm[lm_offset + 42]); \
-    acc.sb += dot(a_reg.s4567, b_lm[lm_offset + 43]); \
-    acc.sc += dot(a_reg.s4567, b_lm[lm_offset + 44]); \
-    acc.sd += dot(a_reg.s4567, b_lm[lm_offset + 45]); \
-    acc.se += dot(a_reg.s4567, b_lm[lm_offset + 46]); \
-    acc.sf += dot(a_reg.s4567, b_lm[lm_offset + 47]); \
-    c_reg.lo += convert_float8(acc.lo); \
-    c_reg.hi += convert_float8(acc.hi); \
-    acc.s0 = dot(a_reg.s89ab, b_lm[lm_offset + 64]); \
-    acc.s1 = dot(a_reg.s89ab, b_lm[lm_offset + 65]); \
-    acc.s2 = dot(a_reg.s89ab, b_lm[lm_offset + 66]); \
-    acc.s3 = dot(a_reg.s89ab, b_lm[lm_offset + 67]); \
-    acc.s4 = dot(a_reg.s89ab, b_lm[lm_offset + 68]); \
-    acc.s5 = dot(a_reg.s89ab, b_lm[lm_offset + 69]); \
-    acc.s6 = dot(a_reg.s89ab, b_lm[lm_offset + 70]); \
-    acc.s7 = dot(a_reg.s89ab, b_lm[lm_offset + 71]); \
-    acc.s8 = dot(a_reg.s89ab, b_lm[lm_offset + 72]); \
-    acc.s9 = dot(a_reg.s89ab, b_lm[lm_offset + 73]); \
-    acc.sa = dot(a_reg.s89ab, b_lm[lm_offset + 74]); \
-    acc.sb = dot(a_reg.s89ab, b_lm[lm_offset + 75]); \
-    acc.sc = dot(a_reg.s89ab, b_lm[lm_offset + 76]); \
-    acc.sd = dot(a_reg.s89ab, b_lm[lm_offset + 77]); \
-    acc.se = dot(a_reg.s89ab, b_lm[lm_offset + 78]); \
-    acc.sf = dot(a_reg.s89ab, b_lm[lm_offset + 79]); \
-    acc.s0 += dot(a_reg.scdef, b_lm[lm_offset + 96]); \
-    acc.s1 += dot(a_reg.scdef, b_lm[lm_offset + 97]); \
-    acc.s2 += dot(a_reg.scdef, b_lm[lm_offset + 98]); \
-    acc.s3 += dot(a_reg.scdef, b_lm[lm_offset + 99]); \
-    acc.s4 += dot(a_reg.scdef, b_lm[lm_offset + 100]); \
-    acc.s5 += dot(a_reg.scdef, b_lm[lm_offset + 101]); \
-    acc.s6 += dot(a_reg.scdef, b_lm[lm_offset + 102]); \
-    acc.s7 += dot(a_reg.scdef, b_lm[lm_offset + 103]); \
-    acc.s8 += dot(a_reg.scdef, b_lm[lm_offset + 104]); \
-    acc.s9 += dot(a_reg.scdef, b_lm[lm_offset + 105]); \
-    acc.sa += dot(a_reg.scdef, b_lm[lm_offset + 106]); \
-    acc.sb += dot(a_reg.scdef, b_lm[lm_offset + 107]); \
-    acc.sc += dot(a_reg.scdef, b_lm[lm_offset + 108]); \
-    acc.sd += dot(a_reg.scdef, b_lm[lm_offset + 109]); \
-    acc.se += dot(a_reg.scdef, b_lm[lm_offset + 110]); \
-    acc.sf += dot(a_reg.scdef, b_lm[lm_offset + 111]); \
-    c_reg.lo += convert_float8(acc.lo); \
-    c_reg.hi += convert_float8(acc.hi); \
-
-
-__attribute__((qcom_wave_pair_mode(1)))
-kernel void kernel_gemm_moe_q8_0_f32_ns(
-        __global     char *           src0_q,   // flat q8_0 quants  [n_expert*ne01*ne00]
-        __global     half *           src0_d,   // flat q8_0 scales  [n_expert*ne01*nb]
-        __read_only  image1d_buffer_t src1,     // reordered activations (f32)
-        __global     uint *           src2,     // post-router out indices
-        __global     ushort *         src2_emap,// expert per tile
-        __write_only image1d_buffer_t dst,
-        __global     int *            total_tiles,
-        uint ne00,
-        uint ne01
-) {
-    uint block_id_m = get_global_id(1); // m_tile
-    uint block_id_n = get_global_id(2); // n_tile
-
-    if (block_id_n >= total_tiles[0]) {
-        return;
-    }
-
-    __private half16 reg_a;
-    __private float32 reg_c = (float32)(0);
-    __local half4 shared_b[128];
-
-    const ushort expert_id = src2_emap[block_id_n];
-
-    const uint row = block_id_m * TILESIZE_M;
-    const uint col = block_id_n * TILESIZE_N;
-
-    const uint nb = ne00 >> 5;                 // blocks per row (ne00/32)
-    const uint w_row = expert_id * ne01 + row + get_local_id(0); // this lane's output row
-    __global char * w_q = src0_q + (ulong)w_row * ne00;          // char base for the row
-    __global half * w_d = src0_d + (ulong)w_row * nb;            // scale base for the row
-
-    uint sub_block_id_m = get_local_id(0);
-    uint2 b_global_offset;
-    b_global_offset.x = ((sub_block_id_m & 3) << 2) + (sub_block_id_m >> 2) * ne00;
-    b_global_offset.y = b_global_offset.x + (16 * ne00);
-    uint2 b_local_offset;
-    b_local_offset.x = (sub_block_id_m & 3) * 32 + (sub_block_id_m >> 2);
-    b_local_offset.y = b_local_offset.x + 16;
-
-    // Loop along K axis, 32 elements per iteration, split into 2 sub-blocks.
-    for (uint step = 0; step < ne00; step += TILESIZE_K * 2) {
-        half s = w_d[step >> 5];               // one q8_0 scale per 32-element block
-
-        // First sub-block: 16 weights (16 chars = one uint4) at K=step
-        uint4 q8x16 = *((__global uint4 *)(w_q + step));
-
-        uint b_sub_offset = col * ne00 + step;
-        float8 bx8_f32;
-        bx8_f32.lo = read_imagef(src1, (b_sub_offset + b_global_offset.x) / 4);
-        bx8_f32.hi = read_imagef(src1, (b_sub_offset + b_global_offset.y) / 4);
-        half8 bx8_f16 = convert_half8(bx8_f32);
-        shared_b[b_local_offset.x] = bx8_f16.lo;
-        shared_b[b_local_offset.y] = bx8_f16.hi;
-
-        dequantize_q8_0(q8x16, reg_a, s);
-
-        sub_group_barrier(CLK_LOCAL_MEM_FENCE);
-
-        half16 acc;
-        dotx16_reduce8(reg_a, shared_b, reg_c.lo, 0);
-        dotx16_reduce8(reg_a, shared_b, reg_c.hi, 16);
-
-        // Second sub-block: next 16 weights at K=step+16
-        uint half_step = step + TILESIZE_K;
-        q8x16 = *((__global uint4 *)(w_q + half_step));
-        b_sub_offset = col * ne00 + half_step;
-
-        bx8_f32.lo = read_imagef(src1, (b_sub_offset + b_global_offset.x) / 4);
-        bx8_f32.hi = read_imagef(src1, (b_sub_offset + b_global_offset.y) / 4);
-        bx8_f16 = convert_half8(bx8_f32);
-        shared_b[b_local_offset.x] = bx8_f16.lo;
-        shared_b[b_local_offset.y] = bx8_f16.hi;
-
-        dequantize_q8_0(q8x16, reg_a, s);
-
-        sub_group_barrier(CLK_LOCAL_MEM_FENCE);
-
-        dotx16_reduce8(reg_a, shared_b, reg_c.lo, 0);
-        dotx16_reduce8(reg_a, shared_b, reg_c.hi, 16);
-    }
-
-    if ((get_global_id(0) + block_id_m * TILESIZE_M) >= ne01) {
-        return;
-    }
-
-    __local uint out_idx[TILESIZE_N];
-
-    if (get_local_id(0) < TILESIZE_N) {
-        uint idx = src2[block_id_n * TILESIZE_N + get_local_id(0)];
-        if (idx == 0xFFFFFFFF) {
-            idx = src2[block_id_n * TILESIZE_N + 0];
-        }
-        out_idx[get_local_id(0)] = idx * ne01;
-    }
-
-    barrier(CLK_LOCAL_MEM_FENCE);
-
-    uint m_offset = row + get_local_id(0);
-
-    write_imagef(dst, out_idx[1] + m_offset, (reg_c.s1));
-    write_imagef(dst, out_idx[2] + m_offset, (reg_c.s2));
-    write_imagef(dst, out_idx[3] + m_offset, (reg_c.s3));
-    write_imagef(dst, out_idx[4] + m_offset, (reg_c.s4));
-    write_imagef(dst, out_idx[5] + m_offset, (reg_c.s5));
-    write_imagef(dst, out_idx[6] + m_offset, (reg_c.s6));
-    write_imagef(dst, out_idx[7] + m_offset, (reg_c.s7));
-    write_imagef(dst, out_idx[8] + m_offset, (reg_c.s8));
-    write_imagef(dst, out_idx[9] + m_offset, (reg_c.s9));
-    write_imagef(dst, out_idx[10] + m_offset, (reg_c.sa));
-    write_imagef(dst, out_idx[11] + m_offset, (reg_c.sb));
-    write_imagef(dst, out_idx[12] + m_offset, (reg_c.sc));
-    write_imagef(dst, out_idx[13] + m_offset, (reg_c.sd));
-    write_imagef(dst, out_idx[14] + m_offset, (reg_c.se));
-    write_imagef(dst, out_idx[15] + m_offset, (reg_c.sf));
-    write_imagef(dst, out_idx[16] + m_offset, (reg_c.sg));
-    write_imagef(dst, out_idx[17] + m_offset, (reg_c.sh));
-    write_imagef(dst, out_idx[18] + m_offset, (reg_c.si));
-    write_imagef(dst, out_idx[19] + m_offset, (reg_c.sj));
-    write_imagef(dst, out_idx[20] + m_offset, (reg_c.sk));
-    write_imagef(dst, out_idx[21] + m_offset, (reg_c.sl));
-    write_imagef(dst, out_idx[22] + m_offset, (reg_c.sm));
-    write_imagef(dst, out_idx[23] + m_offset, (reg_c.sn));
-    write_imagef(dst, out_idx[24] + m_offset, (reg_c.so));
-    write_imagef(dst, out_idx[25] + m_offset, (reg_c.sp));
-    write_imagef(dst, out_idx[26] + m_offset, (reg_c.sq));
-    write_imagef(dst, out_idx[27] + m_offset, (reg_c.sr));
-    write_imagef(dst, out_idx[28] + m_offset, (reg_c.ss));
-    write_imagef(dst, out_idx[29] + m_offset, (reg_c.st));
-    write_imagef(dst, out_idx[30] + m_offset, (reg_c.su));
-    write_imagef(dst, out_idx[31] + m_offset, (reg_c.sv));
-
-    barrier(CLK_GLOBAL_MEM_FENCE);
-    write_imagef(dst, out_idx[0] + m_offset, (reg_c.s0));
-}

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

@@ -1,221 +0,0 @@
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-#pragma OPENCL EXTENSION cl_khr_subgroups : enable
-#ifdef cl_khr_integer_dot_product
-#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
-#endif
-
-// Generic int8 dp4a MoE GEMM, specialized versions also exist
-// MOE_QT:
-//   4 (q4_K)/41(q4_1)/40(q4_0)   NIBBLE   image low nibbles -> EXP4
-//   5 (q5_K)/51(q5_1)/50(q5_0)   NIBBLE+HI image nibbles + qh high-bit plane
-//   6 (q6_K)                     Q6       image nibbles + qh 2-bit -> SIGN6((nibble|hi2))
-//   80(q8_0)/82(mxfp4)           INT8     global int8 codes (mxfp4: convert applies kvalues LUT)
-
-#define TILESIZE_M 64
-#define TILESIZE_N 32
-#define QK_K 256
-
-#ifndef MOE_QT
-#define MOE_QT 4
-#endif
-
-// 4 nibbles in low 16 bits of u -> 4 bytes (value 0..15)
-#define EXP4(u)  ( ((uint)((u) & 0x000Fu))        | \
-                  (((uint)((u) & 0x00F0u)) << 4)  | \
-                  (((uint)((u) & 0x0F00u)) << 8)  | \
-                  (((uint)((u) & 0xF000u)) << 12) )
-// 4 2-bit highs in byte b -> 4 bytes, bits 4-5 (q6_K)
-#define EXP2(b)  ( (((uint)((b) & 0x03u)) << 4)   | \
-                  (((uint)((b) & 0x0Cu)) << 10)  | \
-                  (((uint)((b) & 0x30u)) << 16)  | \
-                  (((uint)((b) & 0xC0u)) << 22) )
-
-// q6 (0..63) -> (q6-32) signed int8/byte (no inter-byte carry)
-inline uint SIGN6(uint q6p){ uint x=q6p^0x20202020u; uint s=x&0x20202020u; return x|(s<<1)|(s<<2); }
-
-// 4 high bits (one per element, in bits 0..3 of h) -> bit4 of each of 4 bytes (5-bit hi)
-#define EXP1(h)  ( (((uint)((h) & 0x1u)) << 4)   | \
-                  (((uint)((h) & 0x2u)) << 11)  | \
-                  (((uint)((h) & 0x4u)) << 18)  | \
-                  (((uint)((h) & 0x8u)) << 25) )
-
-// per-type weight params + per-32-step unpack into qw[8] (8 int8 uints)
-#if MOE_QT == 4 || MOE_QT == 41 || MOE_QT == 40
-  #define WEIGHT_PARAMS __read_only image1d_buffer_t src0_q,
-  #define LOAD_QW(step, sub) \
-      uint qw[8]; { \
-        const uint qoff0 = row + ((ne01*(step))>>3)      + ((expert_id*ne00*ne01)>>3); \
-        const uint qoff1 = row + ((ne01*((step)+16))>>3) + ((expert_id*ne00*ne01)>>3); \
-        const uint r0=read_imageui(src0_q,qoff0+lid).x, r1=read_imageui(src0_q,qoff0+lid+ne01).x; \
-        const uint r2=read_imageui(src0_q,qoff1+lid).x, r3=read_imageui(src0_q,qoff1+lid+ne01).x; \
-        qw[0]=EXP4(r0); qw[1]=EXP4(r0>>16); qw[2]=EXP4(r1); qw[3]=EXP4(r1>>16); \
-        qw[4]=EXP4(r2); qw[5]=EXP4(r2>>16); qw[6]=EXP4(r3); qw[7]=EXP4(r3>>16); }
-
-#elif MOE_QT == 5 || MOE_QT == 51 || MOE_QT == 50
-  // low nibbles via image (q4_K layout) + high-bit plane src0_qh: 1 uint per 32-block
-  // (bit i = high bit of element i). qh laid out [expert][block][row] to match the
-  // existing q5_0 trans4 convert
-  #define WEIGHT_PARAMS __read_only image1d_buffer_t src0_q, __global uint * src0_qh,
-  #define LOAD_QW(step, sub) \
-      uint qw[8]; { \
-        const uint qoff0 = row + ((ne01*(step))>>3)      + ((expert_id*ne00*ne01)>>3); \
-        const uint qoff1 = row + ((ne01*((step)+16))>>3) + ((expert_id*ne00*ne01)>>3); \
-        const uint r0=read_imageui(src0_q,qoff0+lid).x, r1=read_imageui(src0_q,qoff0+lid+ne01).x; \
-        const uint r2=read_imageui(src0_q,qoff1+lid).x, r3=read_imageui(src0_q,qoff1+lid+ne01).x; \
-        const uint h = src0_qh[row_idx + (sub)*ne01 + expert_id*(ne00>>5)*ne01]; \
-        qw[0]=EXP4(r0)|EXP1(h);        qw[1]=EXP4(r0>>16)|EXP1(h>>4); \
-        qw[2]=EXP4(r1)|EXP1(h>>8);     qw[3]=EXP4(r1>>16)|EXP1(h>>12); \
-        qw[4]=EXP4(r2)|EXP1(h>>16);    qw[5]=EXP4(r2>>16)|EXP1(h>>20); \
-        qw[6]=EXP4(r3)|EXP1(h>>24);    qw[7]=EXP4(r3>>16)|EXP1(h>>28); }
-
-#elif MOE_QT == 6
-  #define WEIGHT_PARAMS __read_only image1d_buffer_t src0_ql, __global uint * src0_qh,
-  #define LOAD_QW(step, sub) \
-      uint qw[8]; { \
-        const uint qoff0 = row + ((ne01*(step))>>3)      + ((expert_id*ne00*ne01)>>3); \
-        const uint qoff1 = row + ((ne01*((step)+16))>>3) + ((expert_id*ne00*ne01)>>3); \
-        const uint r0=read_imageui(src0_ql,qoff0+lid).x, r1=read_imageui(src0_ql,qoff0+lid+ne01).x; \
-        const uint r2=read_imageui(src0_ql,qoff1+lid).x, r3=read_imageui(src0_ql,qoff1+lid+ne01).x; \
-        const uint qhb = row + ((sub)*2)*ne01 + expert_id*((ne00>>5)*2)*ne01 + lid; \
-        const uint qh1=src0_qh[qhb], qh2=src0_qh[qhb+ne01]; \
-        qw[0]=SIGN6(EXP4(r0)|EXP2(qh1&0xFFu));        qw[1]=SIGN6(EXP4(r0>>16)|EXP2((qh1>>8)&0xFFu)); \
-        qw[2]=SIGN6(EXP4(r1)|EXP2((qh1>>16)&0xFFu));  qw[3]=SIGN6(EXP4(r1>>16)|EXP2((qh1>>24)&0xFFu)); \
-        qw[4]=SIGN6(EXP4(r2)|EXP2(qh2&0xFFu));        qw[5]=SIGN6(EXP4(r2>>16)|EXP2((qh2>>8)&0xFFu)); \
-        qw[6]=SIGN6(EXP4(r3)|EXP2((qh2>>16)&0xFFu));  qw[7]=SIGN6(EXP4(r3>>16)|EXP2((qh2>>24)&0xFFu)); }
-
-#elif MOE_QT == 80 || MOE_QT == 82
-  // 8-bit direct: int8 codes 8 uints / 32-block, [expert][row][8*sub]. mxfp4: the
-  // convert resolves kvalues_mxfp4[nibble] -> int8 and stores the e8m0_half scale.
-  #define WEIGHT_PARAMS __global uint * src0_q8,
-  #define LOAD_QW(step, sub) \
-      uint qw[8]; { \
-        const uint qb = (expert_id*ne01 + row_idx)*(ne00>>2) + (sub)*8; \
-        qw[0]=src0_q8[qb+0]; qw[1]=src0_q8[qb+1]; qw[2]=src0_q8[qb+2]; qw[3]=src0_q8[qb+3]; \
-        qw[4]=src0_q8[qb+4]; qw[5]=src0_q8[qb+5]; qw[6]=src0_q8[qb+6]; qw[7]=src0_q8[qb+7]; }
-#else
-  #error "unknown MOE_QT"
-#endif
-
-inline int dp4a4(uint w0,uint w1,uint w2,uint w3,uint a0,uint a1,uint a2,uint a3){
-    int r=0; r=dot_acc_sat_4x8packed_ss_int(w0,a0,r); r=dot_acc_sat_4x8packed_ss_int(w1,a1,r);
-    r=dot_acc_sat_4x8packed_ss_int(w2,a2,r); r=dot_acc_sat_4x8packed_ss_int(w3,a3,r); return r; }
-
-// One token's two-half dp4a + uniform scale/min epilogue into acc[t].
-#define MOE_DP4A_T(t) do {                                                                  \
-        const int raw1 = dp4a4(qw[0],qw[1],qw[2],qw[3], sh_qa[t][0],sh_qa[t][1],sh_qa[t][2],sh_qa[t][3]); \
-        const int raw2 = dp4a4(qw[4],qw[5],qw[6],qw[7], sh_qa[t][4],sh_qa[t][5],sh_qa[t][6],sh_qa[t][7]); \
-        const float a_d = (float)sh_d[t];                                                   \
-        acc[t] += sc0*a_d*(float)raw1 + sc1*a_d*(float)raw2 - mn*(float)sh_s[t];             \
-    } while (0)
-
-__attribute__((qcom_wave_pair_mode(1)))
-kernel void kernel_gemm_moe_q8_1_dp4a(
-        WEIGHT_PARAMS                            // per-type native weight buffer(s)
-        __global     half *           src0_scale,// uniform f16 16/superblock (per-16), [expert,row]
-        __global     half *           src0_min,  // uniform f16  8/superblock (per-32), [expert,row]
-        __global     uint *           src1_qa,   // q8_1 activations int8 (as uint, 4/elem)
-        __global     half *           src1_da,   // q8_1 per-block scale [tok_slot * ne00/32]
-        __global     half *           src1_sa,   // q8_1 per-block sum*d [tok_slot * ne00/32]
-        __global     uint *           src2,      // post-router (orig out positions)
-        __global     ushort *         src2_emap, // tile -> expert id
-        __write_only image1d_buffer_t dst,
-        __global     int *            total_tiles,
-        uint ne00,
-        uint ne01,
-        int  is_ragged,
-        int  has_min                             // 0 for symmetric types (q8_0/q6_K/q4_0/...): skip min read
-) {
-    const uint block_id_m = get_global_id(1);
-    const uint block_id_n = get_global_id(2);
-    if (block_id_n >= total_tiles[0]) return;
-
-    const uint lid = get_local_id(0);            // 0..63 -> output row within M-tile
-    const ushort expert_id = src2_emap[block_id_n];
-    const uint   row = block_id_m * TILESIZE_M;
-    const uint   col = block_id_n * TILESIZE_N;
-    const uint   row_idx = row + lid;
-
-    // Scale/min are laid out FLAT per-32-block (2 per-16-segment scales + 1 min per
-    // 32-block), so K only needs to be a multiple of 32 — works for the 32-block
-    // types (q8_0/q5_0/q4_0/...) as well as the K-quants (K%256==0, same bytes).
-    const uint nblk32     = ne00 / 32;
-    const uint sc_per_row = nblk32 * 2;
-    const uint mn_per_row = nblk32;
-    const uint ne00_u = ne00 >> 2;
-    const uint ne00_b = ne00 >> 5;
-
-    __local uint sh_qa[TILESIZE_N][8];
-    __local half sh_d[TILESIZE_N];
-    __local half sh_s[TILESIZE_N];
-
-    __local uint sh_src2[TILESIZE_N];
-    __local int  sh_nreal;
-    if (lid < TILESIZE_N) sh_src2[lid] = src2[col + lid];
-    barrier(CLK_LOCAL_MEM_FENCE);
-    if (lid == 0) {
-        int nr = TILESIZE_N;
-        if (is_ragged) { nr = 0;
-            #pragma unroll
-            for (int t = 0; t < TILESIZE_N; ++t) if (sh_src2[t] != 0xFFFFFFFFu) ++nr; }
-        sh_nreal = nr;
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-    const int n_real = sh_nreal;
-
-    float acc[TILESIZE_N];
-    #pragma unroll
-    for (int t = 0; t < TILESIZE_N; ++t) acc[t] = 0.0f;
-
-    for (uint step = 0; step < ne00; step += 32) {
-        const uint sub = step >> 5;        // 32-block index along K
-
-        // uniform pre-decoded scale (2 per-16-seg) + min (1) for this row, this 32-block
-        __global half * scl = src0_scale + (expert_id*ne01 + row_idx)*sc_per_row + sub*2;
-        const float sc0 = (float)scl[0];
-        const float sc1 = (float)scl[1];
-        float mn = 0.0f;
-        if (has_min) mn = (float)src0_min[(expert_id*ne01 + row_idx)*mn_per_row + sub];
-
-        LOAD_QW(step, sub)
-
-        const uint stage_lim = (uint)n_real * 8;
-        for (uint idx = lid; idx < stage_lim; idx += 64) {
-            const uint t = idx >> 3, u = idx & 7;
-            sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u];
-        }
-        if (lid < (uint)n_real) {
-            sh_d[lid] = src1_da[(col + lid) * ne00_b + sub];
-            sh_s[lid] = src1_sa[(col + lid) * ne00_b + sub];
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-
-        if (n_real == TILESIZE_N) {
-            #pragma unroll
-            for (int t = 0; t < TILESIZE_N; ++t) { MOE_DP4A_T(t); }
-        } else {
-            #pragma unroll 4
-            for (int t = 0; t < n_real; ++t) { MOE_DP4A_T(t); }
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-    }
-
-    if (row_idx >= ne01) return;
-
-    __local uint out_idx[TILESIZE_N];
-    if (lid < TILESIZE_N) {
-        uint idx = sh_src2[lid];
-        if (idx == 0xFFFFFFFF) idx = sh_src2[0];
-        out_idx[lid] = idx * ne01;
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-
-    const uint m_offset = row + lid;
-    if (n_real == TILESIZE_N) {
-        #pragma unroll
-        for (int t = 1; t < TILESIZE_N; ++t) write_imagef(dst, out_idx[t] + m_offset, acc[t]);
-        barrier(CLK_GLOBAL_MEM_FENCE);
-        write_imagef(dst, out_idx[0] + m_offset, acc[0]);
-    } else {
-        for (int t = 0; t < n_real; ++t) write_imagef(dst, out_idx[t] + m_offset, acc[t]);
-    }
-}

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

@@ -1,143 +0,0 @@
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-#pragma OPENCL EXTENSION cl_khr_subgroups : enable
-#ifdef cl_khr_integer_dot_product
-#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
-#endif
-
-// Weight layout, feature-major:
-//   src0_q[row + (k/4)*m]  ushort = 4 nibbles (K = 4*grp .. +3)
-//   src0_d[row + (k/32)*m] half   = per-32-block scale
-
-#define TILESIZE_N 32
-
-// IQ4_NL non-linear codebook as signed int8, packed 4 codes per uint.
-// divergent nibble lookups read a small __constant uint array + shift,
-// never a byte array because byte-indexed __constant loads serialize on Adreno and tank perf
-//   idx 0-3:  -127,-104,-83,-65 = 0x81,0x98,0xAD,0xBF
-//   idx 4-7:  -49,-35,-22,-10   = 0xCF,0xDD,0xEA,0xF6
-//   idx 8-11:  1, 13, 25, 38    = 0x01,0x0D,0x19,0x26
-//   idx 12-15: 53, 69, 89,113   = 0x35,0x45,0x59,0x71
-__constant uint kvalues_iq4nl_i8x4[4] = {
-    0xBFAD9881u, 0xF6EADDCFu, 0x26190D01u, 0x71594535u
-};
-
-// nibble (0..15) -> its codebook byte in the low 8 bits.
-inline uint iq4nl_code(uint n) {
-    return (kvalues_iq4nl_i8x4[n >> 2] >> ((n & 3u) * 8u)) & 0xFFu;
-}
-
-// 4 nibbles in low 16 bits of u -> 4 codebook int8, packed for dp4a.
-inline uint iq4nl_pack(ushort u) {
-    return  iq4nl_code((uint)( u        & 0xF))
-         | (iq4nl_code((uint)((u >>  4) & 0xF)) <<  8)
-         | (iq4nl_code((uint)((u >>  8) & 0xF)) << 16)
-         | (iq4nl_code((uint)((u >> 12) & 0xF)) << 24);
-}
-
-inline int dot8_q8a(uint8 qw, __local const uint * a) {
-    int r = 0;
-    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a[0], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a[1], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a[2], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a[3], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s4, a[4], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s5, a[5], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s6, a[6], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s7, a[7], r);
-    return r;
-}
-
-__attribute__((qcom_wave_pair_mode(1)))
-kernel void kernel_gemm_noshuffle_iq4_nl_q8_1_dp4a(
-        __global const ushort * src0_q,    // IQ4_NL nibbles (4/ushort, feature-major)
-        __global const half   * src0_d,    // per-32-block scale, feature-major
-        __global const uint   * src1_qa,   // q8_1 activations int8 (as uint, 4/elem) [N, K]
-        __global const half   * src1_da,   // q8_1 per-block scale [N, K/32]
-        __global       float  * dst,
-        ulong  offsetd,
-        int    m,                          // output features (rows)
-        int    n_no_padding,               // tokens (cols)
-        int    k                           // K (== ne00)
-) {
-    dst = (global float *)((global char *)dst + offsetd);
-
-    const uint lid = get_local_id(0);          // 0..63 -> row within the M-tile
-    const uint block_id_m = get_global_id(1);
-    const uint block_id_n = get_global_id(2);
-
-    const uint row      = block_id_m * 64 + lid;
-    const uint col_base = block_id_n * TILESIZE_N;
-    const bool row_valid = row < (uint)m;
-    const uint rrow     = row_valid ? row : 0;  // clamp OOB rows; their writes are masked
-
-    const uint k_u = (uint)k >> 2;   // K in uint (int8x4) units
-    const uint k_b = (uint)k >> 5;   // blocks-of-32 along K
-
-    __local uint sh_qa[TILESIZE_N][8];
-    __local half sh_d[TILESIZE_N];
-
-#define NGROUPS (TILESIZE_N / 4)
-    float4 acc[NGROUPS];
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
-
-    for (uint step = 0; step < (uint)k; step += 32) {
-        const uint sub = step >> 5;
-
-        const float d_w = (float)src0_d[rrow + sub * (uint)m];
-
-        // 8 weight uints (32 codebook int8) for this row, this 32-block.
-        const uint qsbase = rrow + (step >> 2) * (uint)m;
-        uint8 qw;
-        qw.s0 = iq4nl_pack(src0_q[qsbase + 0 * m]);
-        qw.s1 = iq4nl_pack(src0_q[qsbase + 1 * m]);
-        qw.s2 = iq4nl_pack(src0_q[qsbase + 2 * m]);
-        qw.s3 = iq4nl_pack(src0_q[qsbase + 3 * m]);
-        qw.s4 = iq4nl_pack(src0_q[qsbase + 4 * m]);
-        qw.s5 = iq4nl_pack(src0_q[qsbase + 5 * m]);
-        qw.s6 = iq4nl_pack(src0_q[qsbase + 6 * m]);
-        qw.s7 = iq4nl_pack(src0_q[qsbase + 7 * m]);
-
-        // cooperatively stage the 32-token x 32-K int8 activations to lm
-        for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            const uint c = col_base + t;
-            sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
-        }
-        if (lid < TILESIZE_N) {
-            const uint c = col_base + lid;
-            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-
-#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
-        #pragma unroll
-        for (int g = 0; g < NGROUPS; ++g) {
-            const int b = g * 4;
-            float4 rf;
-            rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]);  rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
-            rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]);  rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
-            acc[g] += d_w * LD4(sh_d, b) * rf;
-        }
-#undef LD4
-        barrier(CLK_LOCAL_MEM_FENCE);
-    }
-
-    if (!row_valid) {
-        return;
-    }
-
-    // dst is [token, feature] row-major (stride m): dst[col*m + row].
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) {
-        const uint b = (uint)(g * 4);
-        const float4 a = acc[g];
-        const uint c0 = col_base + b;
-        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
-        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
-        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
-        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
-    }
-#undef NGROUPS
-}

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

@@ -1,127 +0,0 @@
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-#pragma OPENCL EXTENSION cl_khr_subgroups : enable
-#ifdef cl_khr_integer_dot_product
-#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
-#endif
-
-#define TILESIZE_N 32
-
-// Expand the 4 nibbles in the low 16 bits of u into 4 bytes (value 0..15),
-// packed for the int8 dp4a. The -8 zero-point is applied via the sum term.
-#define EXP4(u)  ( ((uint)((u) & 0x000Fu))        | \
-                  (((uint)((u) & 0x00F0u)) << 4)  | \
-                  (((uint)((u) & 0x0F00u)) << 8)  | \
-                  (((uint)((u) & 0xF000u)) << 12) )
-
-inline int dot8_q8a(uint8 qw, __local const uint * a) {
-    int r = 0;
-    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a[0], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a[1], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a[2], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a[3], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s4, a[4], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s5, a[5], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s6, a[6], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s7, a[7], r);
-    return r;
-}
-
-__attribute__((qcom_wave_pair_mode(1)))
-kernel void kernel_gemm_noshuffle_q4_0_q8_1_dp4a(
-        __global const ushort * src0_q,    // q4_0 nibbles (4/ushort, feature-major)
-        __global const half   * src0_d,    // per-32-block scale, feature-major
-        __global const uint   * src1_qa,   // q8_1 activations int8 (as uint, 4/elem) [N, K]
-        __global const half   * src1_da,   // q8_1 per-block scale [N, K/32]
-        __global const half   * src1_sa,   // q8_1 per-block sum*d  [N, K/32]
-        __global       float  * dst,
-        ulong  offsetd,
-        int    m,                          // output features (rows)
-        int    n_no_padding,               // tokens (cols)
-        int    k                           // K (== ne00)
-) {
-    dst = (global float *)((global char *)dst + offsetd);
-
-    const uint lid = get_local_id(0);          // 0..63 -> row within the M-tile
-    const uint block_id_m = get_global_id(1);
-    const uint block_id_n = get_global_id(2);
-
-    const uint row      = block_id_m * 64 + lid;
-    const uint col_base = block_id_n * TILESIZE_N;
-    const bool row_valid = row < (uint)m;
-    const uint rrow     = row_valid ? row : 0;  // clamp OOB rows; their writes are masked
-
-    const uint k_u = (uint)k >> 2;   // K in uint (int8x4) units
-    const uint k_b = (uint)k >> 5;   // blocks-of-32 along K
-
-    __local uint sh_qa[TILESIZE_N][8];
-    __local half sh_d[TILESIZE_N];
-    __local half sh_s[TILESIZE_N];
-
-#define NGROUPS (TILESIZE_N / 4)
-    float4 acc[NGROUPS];
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
-
-    for (uint step = 0; step < (uint)k; step += 32) {
-        const uint sub = step >> 5;
-
-        const float d_w = (float)src0_d[rrow + sub * (uint)m];
-
-        // 8 weight uints (32 nibbles) for this row, this 32-block. Feature-major:
-        // src0_q[row + (k/4 + u)*m], k/4 = step/4 (= step>>2). EXP4 -> dp4a int8.
-        const uint qsbase = rrow + (step >> 2) * (uint)m;
-        uint8 qw;
-        qw.s0 = EXP4(src0_q[qsbase + 0 * m]);
-        qw.s1 = EXP4(src0_q[qsbase + 1 * m]);
-        qw.s2 = EXP4(src0_q[qsbase + 2 * m]);
-        qw.s3 = EXP4(src0_q[qsbase + 3 * m]);
-        qw.s4 = EXP4(src0_q[qsbase + 4 * m]);
-        qw.s5 = EXP4(src0_q[qsbase + 5 * m]);
-        qw.s6 = EXP4(src0_q[qsbase + 6 * m]);
-        qw.s7 = EXP4(src0_q[qsbase + 7 * m]);
-
-        // cooperatively stage the 32-token x 32-K int8 activations to LDS
-        for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            const uint c = col_base + t;
-            sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
-        }
-        if (lid < TILESIZE_N) {
-            const uint c = col_base + lid;
-            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
-            sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-
-#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
-        #pragma unroll
-        for (int g = 0; g < NGROUPS; ++g) {
-            const int b = g * 4;
-            float4 rf;
-            rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]);  rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
-            rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]);  rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
-            // q4_0: w = d*(q-8) -> d_w * (a_d * dp4a(q,qa) - 8 * a_s)
-            acc[g] += d_w * (LD4(sh_d, b) * rf - 8.0f * LD4(sh_s, b));
-        }
-#undef LD4
-        barrier(CLK_LOCAL_MEM_FENCE);
-    }
-
-    if (!row_valid) {
-        return;
-    }
-
-    // dst is [token, feature] row-major (stride m): dst[col*m + row].
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) {
-        const uint b = (uint)(g * 4);
-        const float4 a = acc[g];
-        const uint c0 = col_base + b;
-        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
-        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
-        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
-        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
-    }
-#undef NGROUPS
-}

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

@@ -1,281 +0,0 @@
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-#pragma OPENCL EXTENSION cl_khr_subgroups : enable
-#ifdef cl_khr_integer_dot_product
-#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
-#endif
-
-#ifndef TILESIZE_N
-#define TILESIZE_N 32
-#endif
-#define QK_K 256
-#define K_SCALE_SIZE 12
-
-inline void get_scale_min_k4(
-    int j,
-    global const uchar * q,
-    uchar * d,
-    uchar * m,
-    uchar mask_d6,
-    uchar mask_d4,
-    uchar mask_hi2
-) {
-    if (j < 4) {
-        *d = q[j]   & mask_d6;
-        *m = q[j+4] & mask_d6;
-    } else {
-        *d = (q[j+4] & mask_d4) | ((q[j-4] & mask_hi2) >> 2);
-        *m = ((q[j+4] >> 4) & mask_d4) | ((q[j]   & mask_hi2) >> 2);
-    }
-}
-
-// Expand the 4 nibbles in the low 16 bits of `u` into 4 bytes (one nibble per
-// byte, value 0..15), packed for the int8 dp4a.
-#define EXP4(u)  ( ((uint)((u) & 0x000Fu))        | \
-                  (((uint)((u) & 0x00F0u)) << 4)  | \
-                  (((uint)((u) & 0x0F00u)) << 8)  | \
-                  (((uint)((u) & 0xF000u)) << 12) )
-
-// 32-K dp4a dot of one token's int8 activations (8 packed uints in lm) against the
-// row's 8 packed weight uints. qw passed by value as a uint8 (register), not an array.
-inline int dot8_q8a(uint8 qw, __local const uint * a) {
-    int r = 0;
-    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a[0], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a[1], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a[2], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a[3], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s4, a[4], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s5, a[5], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s6, a[6], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s7, a[7], r);
-    return r;
-}
-
-__attribute__((qcom_wave_pair_mode(1)))
-kernel void kernel_gemm_noshuffle_q4_k_q8_1_dp4a(
-        __global const ushort * src0_q,    // q4_K weights (noshuffle, packed nibbles)
-        __global const uchar  * src0_s,    // 6-bit scale/min codes
-        __global const half   * src0_d,    // per-superblock scale
-        __global const half   * src0_dm,   // per-superblock min
-        __global const uint   * src1_qa,   // q8_1 activations int8 (as uint, 4/elem) [N, K]
-        __global const half   * src1_da,   // q8_1 per-block scale [N, K/32]
-        __global const half   * src1_sa,   // q8_1 per-block sum*d [N, K/32]
-        __global       float  * dst,
-        ulong  offsetd,
-        int    m,                          // output features (rows)
-        int    n_no_padding,               // tokens (cols)
-        int    k,                          // K (== ne00)
-        uchar  mask_d6,
-        uchar  mask_d4,
-        uchar  mask_hi2
-) {
-    dst = (global float *)((global char *)dst + offsetd);
-
-    const uint lid = get_local_id(0);          // 0..63 -> row within the M-tile
-    const uint block_id_m = get_global_id(1);
-    const uint block_id_n = get_global_id(2);
-
-    const uint row      = block_id_m * 64 + lid;
-    const uint col_base = block_id_n * TILESIZE_N;
-    const bool row_valid = row < (uint)m;
-    const uint rrow     = row_valid ? row : 0;  // clamp OOB rows; their writes are masked
-
-    const uint num_superblocks = (uint)k / QK_K;
-    const uint k_u = (uint)k >> 2;   // K in uint (int8x4) units
-    const uint k_b = (uint)k >> 5;   // blocks-of-32 along K
-
-    __local uint sh_qa[TILESIZE_N][8];
-    __local half sh_d[TILESIZE_N];
-    __local half sh_s[TILESIZE_N];
-
-    // One float4 vector-register accumulator per group of 4 tokens (NGROUPS = TILESIZE_N/4).
-#define NGROUPS (TILESIZE_N / 4)
-    float4 acc[NGROUPS];
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) { acc[g] = (float4)(0.0f); }
-
-    for (uint step = 0; step < (uint)k; step += 32) {
-        const uint sub     = step >> 5;
-        const uint sb_idx  = step / QK_K;
-        const uint sub_idx = sub & 7;
-
-        // weight scale/min for this WI's row, this subblock
-        const float dd  = (float)src0_d [rrow + sb_idx * m];
-        const float dmm = (float)src0_dm[rrow + sb_idx * m];
-        global const uchar * sc = src0_s + rrow * num_superblocks * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
-        uchar sv, mn;
-        get_scale_min_k4(sub_idx, sc, &sv, &mn, mask_d6, mask_d4, mask_hi2);
-        const float scale = dd  * (float)sv;
-        const float minv  = dmm * (float)mn;
-
-        // repack this row's 32 weight nibbles into 8 dp4a uints. The packed q4_K
-        // layout stores one ushort = 4 consecutive-K nibbles for a row at
-        // src0_q[row + (K_group)*m], K_group = step/4 + u.
-        const uint wbase = rrow + (step >> 2) * (uint)m;
-        uint8 qw;
-        qw.s0 = EXP4(src0_q[wbase + 0 * m]);
-        qw.s1 = EXP4(src0_q[wbase + 1 * m]);
-        qw.s2 = EXP4(src0_q[wbase + 2 * m]);
-        qw.s3 = EXP4(src0_q[wbase + 3 * m]);
-        qw.s4 = EXP4(src0_q[wbase + 4 * m]);
-        qw.s5 = EXP4(src0_q[wbase + 5 * m]);
-        qw.s6 = EXP4(src0_q[wbase + 6 * m]);
-        qw.s7 = EXP4(src0_q[wbase + 7 * m]);
-
-        // cooperatively stage the 32-token x 32-K int8 activations to lm
-        for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            const uint c = col_base + t;
-            sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
-        }
-        if (lid < TILESIZE_N) {
-            const uint c = col_base + lid;
-            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
-            sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-
-#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
-        #pragma unroll
-        for (int g = 0; g < NGROUPS; ++g) {
-            const int b = g * 4;
-            float4 rf;
-            rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]);  rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
-            rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]);  rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
-            acc[g] += scale * LD4(sh_d, b) * rf - minv * LD4(sh_s, b);
-        }
-#undef LD4
-        barrier(CLK_LOCAL_MEM_FENCE);
-    }
-
-    if (!row_valid) {
-        return;
-    }
-
-    // dst is [token, feature] row-major (stride m): dst[col*m + row]. Scatter each
-    // lane with a per-token padding guard (dst is non-contiguous in token).
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) {
-        const uint b = (uint)(g * 4);
-        const float4 a = acc[g];
-        const uint c0 = col_base + b;
-        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
-        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
-        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
-        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
-    }
-#undef NGROUPS
-}
-
-__attribute__((qcom_wave_pair_mode(1)))
-kernel void kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg(
-        __read_only image1d_buffer_t src0_q_img, // q4_K weights as uint32 texels (2 ushorts/texel)
-        __global const uchar  * src0_s,    // 6-bit scale/min codes
-        __global const half   * src0_d,    // per-superblock scale
-        __global const half   * src0_dm,   // per-superblock min
-        __global const uint   * src1_qa,   // q8_1 activations int8 (as uint, 4/elem) [N, K]
-        __global const half   * src1_da,   // q8_1 per-block scale [N, K/32]
-        __global const half   * src1_sa,   // q8_1 per-block sum*d [N, K/32]
-        __global       float  * dst,
-        ulong  offsetd,
-        int    m,                          // output features (rows)
-        int    n_no_padding,               // tokens (cols)
-        int    k,                          // K (== ne00)
-        uchar  mask_d6,
-        uchar  mask_d4,
-        uchar  mask_hi2
-) {
-    dst = (global float *)((global char *)dst + offsetd);
-
-    const uint lid = get_local_id(0);          // 0..63 -> row within the M-tile
-    const uint block_id_m = get_global_id(1);
-    const uint block_id_n = get_global_id(2);
-
-    const uint row      = block_id_m * 64 + lid;
-    const uint col_base = block_id_n * TILESIZE_N;
-    const bool row_valid = row < (uint)m;
-    const uint rrow     = row_valid ? row : 0;  // clamp OOB rows; their writes are masked
-
-    // Constant per WI: the ushort the row needs always sits in the same half of
-    // its uint32 texel (m even => index parity == rrow parity). Hoist the shift.
-    const uint sel = (rrow & 1u) * 16u;
-
-    const uint k_u = (uint)k >> 2;   // K in uint (int8x4) units
-    const uint k_b = (uint)k >> 5;   // blocks-of-32 along K
-    const uint num_superblocks = (uint)k / QK_K;
-
-    __local uint sh_qa[TILESIZE_N][8];
-    __local half sh_d[TILESIZE_N];
-    __local half sh_s[TILESIZE_N];
-
-#define NGROUPS (TILESIZE_N / 4)
-    float4 acc[NGROUPS];
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
-
-    for (uint step = 0; step < (uint)k; step += 32) {
-        const uint sub     = step >> 5;
-        const uint sb_idx  = step / QK_K;
-        const uint sub_idx = sub & 7;
-
-        const float dd  = (float)src0_d [rrow + sb_idx * m];
-        const float dmm = (float)src0_dm[rrow + sb_idx * m];
-        global const uchar * sc = src0_s + rrow * num_superblocks * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
-        uchar sv, mn;
-        get_scale_min_k4(sub_idx, sc, &sv, &mn, mask_d6, mask_d4, mask_hi2);
-        const float scale = dd  * (float)sv;
-        const float minv  = dmm * (float)mn;
-
-        const uint wbase = rrow + (step >> 2) * (uint)m;
-        uint8 qw;
-        qw.s0 = EXP4(read_imageui(src0_q_img, (int)((wbase + 0 * m) >> 1)).x >> sel);
-        qw.s1 = EXP4(read_imageui(src0_q_img, (int)((wbase + 1 * m) >> 1)).x >> sel);
-        qw.s2 = EXP4(read_imageui(src0_q_img, (int)((wbase + 2 * m) >> 1)).x >> sel);
-        qw.s3 = EXP4(read_imageui(src0_q_img, (int)((wbase + 3 * m) >> 1)).x >> sel);
-        qw.s4 = EXP4(read_imageui(src0_q_img, (int)((wbase + 4 * m) >> 1)).x >> sel);
-        qw.s5 = EXP4(read_imageui(src0_q_img, (int)((wbase + 5 * m) >> 1)).x >> sel);
-        qw.s6 = EXP4(read_imageui(src0_q_img, (int)((wbase + 6 * m) >> 1)).x >> sel);
-        qw.s7 = EXP4(read_imageui(src0_q_img, (int)((wbase + 7 * m) >> 1)).x >> sel);
-
-        for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            const uint c = col_base + t;
-            sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
-        }
-        if (lid < TILESIZE_N) {
-            const uint c = col_base + lid;
-            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
-            sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-
-#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
-        #pragma unroll
-        for (int g = 0; g < NGROUPS; ++g) {
-            const int b = g * 4;
-            float4 rf;
-            rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]);  rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
-            rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]);  rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
-            acc[g] += scale * LD4(sh_d, b) * rf - minv * LD4(sh_s, b);
-        }
-#undef LD4
-        barrier(CLK_LOCAL_MEM_FENCE);
-    }
-
-    if (!row_valid) {
-        return;
-    }
-
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) {
-        const uint b = (uint)(g * 4);
-        const float4 a = acc[g];
-        const uint c0 = col_base + b;
-        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
-        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
-        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
-        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
-    }
-#undef NGROUPS
-}

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

@@ -1,235 +0,0 @@
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-#pragma OPENCL EXTENSION cl_khr_subgroups : enable
-#ifdef cl_khr_integer_dot_product
-#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
-#endif
-
-// Weight layout
-//   src0_qs[row + (k/4)*m]  ushort = 4 low nibbles (K = 4*grp .. +3)
-//   src0_qh[row + (k/8)*m]  uchar  = 8 high bits  (one per element)
-//   src0_d [row + (k/32)*m] half   = per-32-block scale
-
-#define TILESIZE_N 32
-
-// 4 nibbles in low 16 bits of u -> 4 bytes (value 0..15)
-#define EXP4(u)  ( ((uint)((u) & 0x000Fu))        | \
-                  (((uint)((u) & 0x00F0u)) << 4)  | \
-                  (((uint)((u) & 0x0F00u)) << 8)  | \
-                  (((uint)((u) & 0xF000u)) << 12) )
-// 4 high bits (one per element, in bits 0..3 of h) -> bit4 of each of 4 bytes
-#define EXP1(h)  ( (((uint)((h) & 0x1u)) << 4)   | \
-                  (((uint)((h) & 0x2u)) << 11)  | \
-                  (((uint)((h) & 0x4u)) << 18)  | \
-                  (((uint)((h) & 0x8u)) << 25) )
-
-inline int dot8_q8a(uint8 qw, __local const uint * a) {
-    int r = 0;
-    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a[0], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a[1], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a[2], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a[3], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s4, a[4], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s5, a[5], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s6, a[6], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s7, a[7], r);
-    return r;
-}
-
-__attribute__((qcom_wave_pair_mode(1)))
-kernel void kernel_gemm_noshuffle_q5_0_q8_1_dp4a(
-        __global const ushort * src0_qs,   // q5_0 low nibbles (4/ushort, feature-major)
-        __global const uchar  * src0_qh,   // q5_0 high-bit plane (8/uchar, feature-major)
-        __global const half   * src0_d,    // per-32-block scale, feature-major
-        __global const uint   * src1_qa,   // q8_1 activations int8 (as uint, 4/elem) [N, K]
-        __global const half   * src1_da,   // q8_1 per-block scale [N, K/32]
-        __global const half   * src1_sa,   // q8_1 per-block sum*d [N, K/32]
-        __global       float  * dst,
-        ulong  offsetd,
-        int    m,                          // output features (rows)
-        int    n_no_padding,               // tokens (cols)
-        int    k                           // K (== ne00)
-) {
-    dst = (global float *)((global char *)dst + offsetd);
-
-    const uint lid = get_local_id(0);          // 0..63 -> row within the M-tile
-    const uint block_id_m = get_global_id(1);
-    const uint block_id_n = get_global_id(2);
-
-    const uint row      = block_id_m * 64 + lid;
-    const uint col_base = block_id_n * TILESIZE_N;
-    const bool row_valid = row < (uint)m;
-    const uint rrow     = row_valid ? row : 0;  // clamp OOB rows; their writes are masked
-
-    const uint k_u = (uint)k >> 2;   // K in uint (int8x4) units
-    const uint k_b = (uint)k >> 5;   // blocks-of-32 along K
-
-    __local uint sh_qa[TILESIZE_N][8];
-    __local half sh_d[TILESIZE_N];
-    __local half sh_s[TILESIZE_N];
-
-#define NGROUPS (TILESIZE_N / 4)
-    float4 acc[NGROUPS];
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
-
-    for (uint step = 0; step < (uint)k; step += 32) {
-        const uint sub = step >> 5;
-
-        const float d_w  = (float)src0_d[rrow + sub * (uint)m];
-        const float minv = d_w * 16.0f;     // -16 centering -> subtract via q8_1 sum
-
-        // 8 weight uints (32 elements) for this row, this 32-block.
-        // nibbles: src0_qs[row + (step/4 + u)*m]; high bits: src0_qh[row + (step/8 + u/2)*m],
-        // 4-bit group selected by (u&1)*4.
-        const uint qsbase = rrow + (step >> 2) * (uint)m;
-        const uint qhbase = rrow + (step >> 3) * (uint)m;
-        uint8 qw;
-        #define QW(u) (EXP4(src0_qs[qsbase + (u) * m]) | \
-                       EXP1((uint)(src0_qh[qhbase + ((u) >> 1) * m] >> (((u) & 1u) * 4u)) & 0xFu))
-        qw.s0 = QW(0); qw.s1 = QW(1); qw.s2 = QW(2); qw.s3 = QW(3);
-        qw.s4 = QW(4); qw.s5 = QW(5); qw.s6 = QW(6); qw.s7 = QW(7);
-        #undef QW
-
-        // cooperatively stage the 32-token x 32-K int8 activations to lm
-        for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            const uint c = col_base + t;
-            sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
-        }
-        if (lid < TILESIZE_N) {
-            const uint c = col_base + lid;
-            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
-            sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-
-#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
-        #pragma unroll
-        for (int g = 0; g < NGROUPS; ++g) {
-            const int b = g * 4;
-            float4 rf;
-            rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]);  rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
-            rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]);  rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
-            acc[g] += d_w * LD4(sh_d, b) * rf - minv * LD4(sh_s, b);
-        }
-#undef LD4
-        barrier(CLK_LOCAL_MEM_FENCE);
-    }
-
-    if (!row_valid) {
-        return;
-    }
-
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) {
-        const uint b = (uint)(g * 4);
-        const float4 a = acc[g];
-        const uint c0 = col_base + b;
-        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
-        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
-        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
-        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
-    }
-#undef NGROUPS
-}
-
-__attribute__((qcom_wave_pair_mode(1)))
-kernel void kernel_gemm_noshuffle_q5_0_q8_1_dp4a_wimg(
-        __read_only image1d_buffer_t src0_qs_img, // q5_0 low nibbles as uint32 texels (2 ushorts/texel)
-        __global const uchar  * src0_qh,
-        __global const half   * src0_d,
-        __global const uint   * src1_qa,
-        __global const half   * src1_da,
-        __global const half   * src1_sa,
-        __global       float  * dst,
-        ulong  offsetd,
-        int    m,
-        int    n_no_padding,
-        int    k
-) {
-    dst = (global float *)((global char *)dst + offsetd);
-
-    const uint lid = get_local_id(0);
-    const uint block_id_m = get_global_id(1);
-    const uint block_id_n = get_global_id(2);
-
-    const uint row      = block_id_m * 64 + lid;
-    const uint col_base = block_id_n * TILESIZE_N;
-    const bool row_valid = row < (uint)m;
-    const uint rrow     = row_valid ? row : 0;
-
-    const uint sel = (rrow & 1u) * 16u;   // constant per WI: qs ushort half in its uint32 texel
-
-    const uint k_u = (uint)k >> 2;
-    const uint k_b = (uint)k >> 5;
-
-    __local uint sh_qa[TILESIZE_N][8];
-    __local half sh_d[TILESIZE_N];
-    __local half sh_s[TILESIZE_N];
-
-#define NGROUPS (TILESIZE_N / 4)
-    float4 acc[NGROUPS];
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
-
-    for (uint step = 0; step < (uint)k; step += 32) {
-        const uint sub = step >> 5;
-
-        const float d_w  = (float)src0_d[rrow + sub * (uint)m];
-        const float minv = d_w * 16.0f;
-
-        const uint qsbase = rrow + (step >> 2) * (uint)m;   // ushort index
-        const uint qhbase = rrow + (step >> 3) * (uint)m;
-        uint8 qw;
-        // qs ushort via texture: uint32 texel = ushort_index>>1, half = sel.
-        #define QSU(u) ((read_imageui(src0_qs_img, (int)((qsbase + (u) * m) >> 1)).x >> sel) & 0xFFFFu)
-        #define QW(u) (EXP4(QSU(u)) | \
-                       EXP1((uint)(src0_qh[qhbase + ((u) >> 1) * m] >> (((u) & 1u) * 4u)) & 0xFu))
-        qw.s0 = QW(0); qw.s1 = QW(1); qw.s2 = QW(2); qw.s3 = QW(3);
-        qw.s4 = QW(4); qw.s5 = QW(5); qw.s6 = QW(6); qw.s7 = QW(7);
-        #undef QW
-        #undef QSU
-
-        for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            const uint c = col_base + t;
-            sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
-        }
-        if (lid < TILESIZE_N) {
-            const uint c = col_base + lid;
-            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
-            sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-
-#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
-        #pragma unroll
-        for (int g = 0; g < NGROUPS; ++g) {
-            const int b = g * 4;
-            float4 rf;
-            rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]);  rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
-            rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]);  rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
-            acc[g] += d_w * LD4(sh_d, b) * rf - minv * LD4(sh_s, b);
-        }
-#undef LD4
-        barrier(CLK_LOCAL_MEM_FENCE);
-    }
-
-    if (!row_valid) {
-        return;
-    }
-
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) {
-        const uint b = (uint)(g * 4);
-        const float4 a = acc[g];
-        const uint c0 = col_base + b;
-        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
-        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
-        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
-        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
-    }
-#undef NGROUPS
-}

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

@@ -1,164 +0,0 @@
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-#pragma OPENCL EXTENSION cl_khr_subgroups : enable
-#ifdef cl_khr_integer_dot_product
-#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
-#endif
-
-#define TILESIZE_N 32
-#define QK_K 256
-#define K_SCALE_SIZE 12
-
-inline void get_scale_min_k4(
-    int j,
-    global const uchar * q,
-    uchar * d,
-    uchar * m,
-    uchar mask_d6,
-    uchar mask_d4,
-    uchar mask_hi2
-) {
-    if (j < 4) {
-        *d = q[j]   & mask_d6;
-        *m = q[j+4] & mask_d6;
-    } else {
-        *d = (q[j+4] & mask_d4) | ((q[j-4] & mask_hi2) >> 2);
-        *m = ((q[j+4] >> 4) & mask_d4) | ((q[j]   & mask_hi2) >> 2);
-    }
-}
-
-// 4 nibbles in the low 16 bits of `u` -> 4 bytes (value 0..15, bits 0-3).
-#define EXP4(u)  ( ((uint)((u) & 0x000Fu))        | \
-                  (((uint)((u) & 0x00F0u)) << 4)  | \
-                  (((uint)((u) & 0x0F00u)) << 8)  | \
-                  (((uint)((u) & 0xF000u)) << 12) )
-
-// 4 high bits (one per element, in bits 0-3 of h) -> bit 4 of each of 4 bytes,
-// so OR with EXP4 forms the 5-bit q5_K code 0..31.
-#define EXP1(h)  ( (((uint)((h) & 0x1u)) << 4)   | \
-                  (((uint)((h) & 0x2u)) << 11)  | \
-                  (((uint)((h) & 0x4u)) << 18)  | \
-                  (((uint)((h) & 0x8u)) << 25) )
-
-inline int dot8_q8a(uint8 qw, __local const uint * a) {
-    int r = 0;
-    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a[0], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a[1], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a[2], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a[3], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s4, a[4], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s5, a[5], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s6, a[6], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s7, a[7], r);
-    return r;
-}
-
-__attribute__((qcom_wave_pair_mode(1)))
-kernel void kernel_gemm_noshuffle_q5_k_q8_1_dp4a(
-        __global const ushort * src0_q,    // q5_K low nibbles (transposed, ushort = 4 nibbles)
-        __global const uchar  * src0_qh,   // q5_K high bits (transposed, uchar = 8 elems/byte)
-        __global const uchar  * src0_s,    // 6-bit scale/min codes [row][superblock][12]
-        __global const half   * src0_d,    // per-superblock scale (transposed)
-        __global const half   * src0_dm,   // per-superblock min (transposed)
-        __global const uint   * src1_qa,   // q8_1 activations int8 (as uint, 4/elem) [N, K]
-        __global const half   * src1_da,   // q8_1 per-block scale [N, K/32]
-        __global const half   * src1_sa,   // q8_1 per-block sum*d [N, K/32]
-        __global       float  * dst,
-        ulong  offsetd,
-        int    m,                          // output features (rows)
-        int    n_no_padding,               // tokens (cols)
-        int    k,                          // K (== ne00)
-        uchar  mask_d6,
-        uchar  mask_d4,
-        uchar  mask_hi2
-) {
-    dst = (global float *)((global char *)dst + offsetd);
-
-    const uint lid = get_local_id(0);          // 0..63 -> row within the M-tile
-    const uint block_id_m = get_global_id(1);
-    const uint block_id_n = get_global_id(2);
-
-    const uint row      = block_id_m * 64 + lid;
-    const uint col_base = block_id_n * TILESIZE_N;
-    const bool row_valid = row < (uint)m;
-    const uint rrow     = row_valid ? row : 0;
-
-    const uint num_superblocks = (uint)k / QK_K;
-    const uint k_u = (uint)k >> 2;
-    const uint k_b = (uint)k >> 5;
-
-    __local uint sh_qa[TILESIZE_N][8];
-    __local half sh_d[TILESIZE_N];
-    __local half sh_s[TILESIZE_N];
-
-#define NGROUPS (TILESIZE_N / 4)
-    float4 acc[NGROUPS];
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
-
-    for (uint step = 0; step < (uint)k; step += 32) {
-        const uint sub     = step >> 5;
-        const uint sb_idx  = step / QK_K;
-        const uint sub_idx = sub & 7;
-
-        const float dd  = (float)src0_d [rrow + sb_idx * m];
-        const float dmm = (float)src0_dm[rrow + sb_idx * m];
-        global const uchar * sc = src0_s + rrow * num_superblocks * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
-        uchar sv, mn;
-        get_scale_min_k4(sub_idx, sc, &sv, &mn, mask_d6, mask_d4, mask_hi2);
-        const float scale = dd  * (float)sv;
-        const float minv  = dmm * (float)mn;
-
-        // repack this row's 32 weights (nibble | high-bit) into 8 dp4a uints.
-        // ushort u -> 4 elements at K = step + u*4; its 4 high bits are nibble
-        // (u&1) of qh byte (step/8 + u/2).
-        const uint wbase  = rrow + (step >> 2) * (uint)m;
-        const uint qhbase = rrow + (step >> 3) * (uint)m;
-        uint8 qw;
-#define QWU(u) ( EXP4((uint)src0_q[wbase + (uint)(u) * m]) \
-               | EXP1( (uint)((src0_qh[qhbase + (uint)((u) >> 1) * m] >> (((u) & 1) * 4)) & 0x0Fu) ) )
-        qw.s0 = QWU(0); qw.s1 = QWU(1); qw.s2 = QWU(2); qw.s3 = QWU(3);
-        qw.s4 = QWU(4); qw.s5 = QWU(5); qw.s6 = QWU(6); qw.s7 = QWU(7);
-#undef QWU
-
-        for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            const uint c = col_base + t;
-            sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
-        }
-        if (lid < TILESIZE_N) {
-            const uint c = col_base + lid;
-            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
-            sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-
-#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
-        #pragma unroll
-        for (int g = 0; g < NGROUPS; ++g) {
-            const int b = g * 4;
-            float4 rf;
-            rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]);  rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
-            rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]);  rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
-            acc[g] += scale * LD4(sh_d, b) * rf - minv * LD4(sh_s, b);
-        }
-#undef LD4
-        barrier(CLK_LOCAL_MEM_FENCE);
-    }
-
-    if (!row_valid) {
-        return;
-    }
-
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) {
-        const uint b = (uint)(g * 4);
-        const float4 a = acc[g];
-        const uint c0 = col_base + b;
-        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
-        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
-        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
-        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
-    }
-#undef NGROUPS
-}

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

@@ -1,144 +0,0 @@
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-#pragma OPENCL EXTENSION cl_khr_subgroups : enable
-#ifdef cl_khr_integer_dot_product
-#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
-#endif
-
-#define TILESIZE_N 32
-#define QK_K 256
-
-// 4 nibbles in the low 16 bits of `u` -> 4 bytes (value 0..15, in bits 0-3).
-#define EXP4(u)  ( ((uint)((u) & 0x000Fu))        | \
-                  (((uint)((u) & 0x00F0u)) << 4)  | \
-                  (((uint)((u) & 0x0F00u)) << 8)  | \
-                  (((uint)((u) & 0xF000u)) << 12) )
-
-// 4 2-bit highs in byte `b` -> 4 bytes, value 0..3 in bits 4-5 (pre-multiplied
-// by 16 so it ORs with the EXP4 nibble to form q6 in 0..63).
-#define EXP2(b)  ( (((uint)((b) & 0x03u)) << 4)   | \
-                  (((uint)((b) & 0x0Cu)) << 10)  | \
-                  (((uint)((b) & 0x30u)) << 16)  | \
-                  (((uint)((b) & 0xC0u)) << 22) )
-
-// q6 (0..63, bits 0-5 of each byte) -> (q6-32) as a signed int8 per byte.
-inline uint SIGN6(uint q6p) {
-    uint x = q6p ^ 0x20202020u;
-    uint s = x & 0x20202020u;
-    return x | (s << 1) | (s << 2);
-}
-
-// 16-K dp4a dot: 4 packed weight uints against 4 packed int8 activation uints.
-inline int dot4_q8a(uint w0, uint w1, uint w2, uint w3,
-                    uint a0, uint a1, uint a2, uint a3) {
-    int r = 0;
-    r = dot_acc_sat_4x8packed_ss_int(w0, a0, r);
-    r = dot_acc_sat_4x8packed_ss_int(w1, a1, r);
-    r = dot_acc_sat_4x8packed_ss_int(w2, a2, r);
-    r = dot_acc_sat_4x8packed_ss_int(w3, a3, r);
-    return r;
-}
-
-__attribute__((qcom_wave_pair_mode(1)))
-kernel void kernel_gemm_noshuffle_q6_k_q8_1_dp4a(
-        __global const ushort * src0_ql,   // q6_K low nibbles (noshuffle)
-        __global const uchar  * src0_qh,   // q6_K high 2-bit (uchar, 4 highs/elem)
-        __global const ushort * src0_s,    // int8 scale codes (2 chars/ushort, per 16)
-        __global const half   * src0_d,    // per-superblock scale
-        __global const uint   * src1_qa,   // q8_1 activations int8 (as uint, 4/elem) [N, K]
-        __global const half   * src1_da,   // q8_1 per-block scale [N, K/32]
-        __global       float  * dst,
-        ulong  offsetd,
-        int    m,                          // output features (rows)
-        int    n_no_padding,               // tokens (cols)
-        int    k                           // K (== ne00)
-) {
-    dst = (global float *)((global char *)dst + offsetd);
-
-    const uint lid = get_local_id(0);          // 0..63 -> row within the M-tile
-    const uint block_id_m = get_global_id(1);
-    const uint block_id_n = get_global_id(2);
-
-    const uint row      = block_id_m * 64 + lid;
-    const uint col_base = block_id_n * TILESIZE_N;
-    const bool row_valid = row < (uint)m;
-    const uint rrow     = row_valid ? row : 0;  // clamp OOB rows; their writes are masked
-
-    const uint k_u = (uint)k >> 2;   // K in uint (int8x4) units
-    const uint k_b = (uint)k >> 5;   // blocks-of-32 along K
-
-    __local uint sh_qa[TILESIZE_N][8];
-    __local half sh_d[TILESIZE_N];
-
-#define NGROUPS (TILESIZE_N / 4)
-    float4 acc[NGROUPS];
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
-
-    for (uint step = 0; step < (uint)k; step += 32) {
-        const uint sub    = step >> 5;    // 32-block index along K
-        const uint sb_idx = step / QK_K;  // superblock index
-
-        // q6_K superblock scale + the two int8 sub-scales spanning this 32-block
-        const float dd = (float)src0_d[rrow + sb_idx * m];
-        const char2 sc = as_char2(src0_s[rrow + sub * m]);
-        const float scale0 = dd * (float)sc.s0;   // K step..step+15
-        const float scale1 = dd * (float)sc.s1;   // K step+16..step+31
-
-        // repack this row's 32 weights into 8 dp4a uints (4 K each). ql ushort +
-        // qh uchar are co-located at src0_*[row + (step/4 + u)*m].
-        const uint wbase = rrow + (step >> 2) * (uint)m;
-        uint qw[8];
-        #pragma unroll
-        for (int u = 0; u < 8; ++u) {
-            const uint o  = wbase + (uint)u * (uint)m;
-            qw[u] = SIGN6(EXP4((uint)src0_ql[o]) | EXP2((uint)src0_qh[o]));
-        }
-
-        // cooperatively stage the 32-token x 32-K int8 activations + scale
-        for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            const uint c = col_base + t;
-            sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
-        }
-        if (lid < TILESIZE_N) {
-            const uint c = col_base + lid;
-            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-
-        #pragma unroll
-        for (int g = 0; g < NGROUPS; ++g) {
-            const int b = g * 4;
-            float4 rf;
-            #define DOT_TOK(j) { \
-                __local const uint * a = sh_qa[b + (j)]; \
-                const int raw1 = dot4_q8a(qw[0], qw[1], qw[2], qw[3], a[0], a[1], a[2], a[3]); \
-                const int raw2 = dot4_q8a(qw[4], qw[5], qw[6], qw[7], a[4], a[5], a[6], a[7]); \
-                rf.s##j = scale0 * (float)raw1 + scale1 * (float)raw2; \
-            }
-            DOT_TOK(0); DOT_TOK(1); DOT_TOK(2); DOT_TOK(3);
-            #undef DOT_TOK
-            const float4 ad = (float4)((float)sh_d[b+0], (float)sh_d[b+1], (float)sh_d[b+2], (float)sh_d[b+3]);
-            acc[g] += ad * rf;
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-    }
-
-    if (!row_valid) {
-        return;
-    }
-
-    // dst is [token, feature] row-major (stride m): dst[col*m + row].
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) {
-        const uint b = (uint)(g * 4);
-        const float4 a = acc[g];
-        const uint c0 = col_base + b;
-        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
-        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
-        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
-        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
-    }
-#undef NGROUPS
-}

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

@@ -1,212 +0,0 @@
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-#pragma OPENCL EXTENSION cl_khr_subgroups : enable
-#ifdef cl_khr_integer_dot_product
-#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
-#endif
-
-// ne1<=8 keeps the f16 / bin small-batch path.
-
-#define TILESIZE_N 32
-
-// 32-K dp4a dot of one token's int8 activations (8 packed uints in lm) against
-// 8 packed weight uints. q8_0 weights are already dp4a-format signed int8.
-inline int dot8_q8a(uint8 qw, __local const uint * a) {
-    int r = 0;
-    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a[0], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a[1], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a[2], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a[3], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s4, a[4], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s5, a[5], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s6, a[6], r);
-    r = dot_acc_sat_4x8packed_ss_int(qw.s7, a[7], r);
-    return r;
-}
-
-__attribute__((qcom_wave_pair_mode(1)))
-kernel void kernel_gemm_noshuffle_q8_0_q8_1_dp4a(
-        __global const uint  * src0_q,     // q8_0 weights: signed int8, 4/uint, feature-major
-        __global const half  * src0_d,     // per-32-block scale, feature-major [row + (k/32)*m]
-        __global const uint  * src1_qa,    // q8_1 activations int8 (as uint, 4/elem) [N, K]
-        __global const half  * src1_da,    // q8_1 per-block scale [N, K/32]
-        __global       float * dst,
-        ulong  offsetd,
-        int    m,                          // output features (rows)
-        int    n_no_padding,               // tokens (cols)
-        int    k                           // K (== ne00)
-) {
-    dst = (global float *)((global char *)dst + offsetd);
-
-    const uint lid = get_local_id(0);          // 0..63 -> row within the M-tile
-    const uint block_id_m = get_global_id(1);
-    const uint block_id_n = get_global_id(2);
-
-    const uint row      = block_id_m * 64 + lid;
-    const uint col_base = block_id_n * TILESIZE_N;
-    const bool row_valid = row < (uint)m;
-    const uint rrow     = row_valid ? row : 0;  // clamp OOB rows; their writes are masked
-
-    const uint k_u = (uint)k >> 2;   // K in uint (int8x4) units
-    const uint k_b = (uint)k >> 5;   // blocks-of-32 along K
-
-    __local uint sh_qa[TILESIZE_N][8];
-    __local half sh_d[TILESIZE_N];
-
-#define NGROUPS (TILESIZE_N / 4)
-    float4 acc[NGROUPS];
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
-
-    for (uint step = 0; step < (uint)k; step += 32) {
-        const uint sub = step >> 5;
-
-        const float d_w = (float)src0_d[rrow + sub * (uint)m];
-
-        // 8 weight uints (32 int8) for this row, this 32-block. Feature-major:
-        // src0_q[row + (k/4 + u)*m], k/4 = step/4 (= step>>2).
-        const uint wbase = rrow + (step >> 2) * (uint)m;
-        uint8 qw;
-        qw.s0 = src0_q[wbase + 0 * m];
-        qw.s1 = src0_q[wbase + 1 * m];
-        qw.s2 = src0_q[wbase + 2 * m];
-        qw.s3 = src0_q[wbase + 3 * m];
-        qw.s4 = src0_q[wbase + 4 * m];
-        qw.s5 = src0_q[wbase + 5 * m];
-        qw.s6 = src0_q[wbase + 6 * m];
-        qw.s7 = src0_q[wbase + 7 * m];
-
-        // cooperatively stage the 32-token x 32-K int8 activations to LDS
-        for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            const uint c = col_base + t;
-            sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
-        }
-        if (lid < TILESIZE_N) {
-            const uint c = col_base + lid;
-            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-
-#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
-        #pragma unroll
-        for (int g = 0; g < NGROUPS; ++g) {
-            const int b = g * 4;
-            float4 rf;
-            rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]);  rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
-            rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]);  rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
-            acc[g] += d_w * LD4(sh_d, b) * rf;
-        }
-#undef LD4
-        barrier(CLK_LOCAL_MEM_FENCE);
-    }
-
-    if (!row_valid) {
-        return;
-    }
-
-    // dst is [token, feature] row-major (stride m): dst[col*m + row].
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) {
-        const uint b = (uint)(g * 4);
-        const float4 a = acc[g];
-        const uint c0 = col_base + b;
-        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
-        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
-        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
-        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
-    }
-#undef NGROUPS
-}
-
-__attribute__((qcom_wave_pair_mode(1)))
-kernel void kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg(
-        __read_only image1d_buffer_t src0_q_img,  // q8_0 weights as uint32 texels (4 int8/texel)
-        __global const half  * src0_d,
-        __global const uint  * src1_qa,
-        __global const half  * src1_da,
-        __global       float * dst,
-        ulong  offsetd,
-        int    m,
-        int    n_no_padding,
-        int    k
-) {
-    dst = (global float *)((global char *)dst + offsetd);
-
-    const uint lid = get_local_id(0);
-    const uint block_id_m = get_global_id(1);
-    const uint block_id_n = get_global_id(2);
-
-    const uint row      = block_id_m * 64 + lid;
-    const uint col_base = block_id_n * TILESIZE_N;
-    const bool row_valid = row < (uint)m;
-    const uint rrow     = row_valid ? row : 0;
-
-    const uint k_u = (uint)k >> 2;
-    const uint k_b = (uint)k >> 5;
-
-    __local uint sh_qa[TILESIZE_N][8];
-    __local half sh_d[TILESIZE_N];
-
-#define NGROUPS (TILESIZE_N / 4)
-    float4 acc[NGROUPS];
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
-
-    for (uint step = 0; step < (uint)k; step += 32) {
-        const uint sub = step >> 5;
-
-        const float d_w = (float)src0_d[rrow + sub * (uint)m];
-
-        const uint wbase = rrow + (step >> 2) * (uint)m;
-        uint8 qw;
-        qw.s0 = read_imageui(src0_q_img, (int)(wbase + 0 * m)).x;
-        qw.s1 = read_imageui(src0_q_img, (int)(wbase + 1 * m)).x;
-        qw.s2 = read_imageui(src0_q_img, (int)(wbase + 2 * m)).x;
-        qw.s3 = read_imageui(src0_q_img, (int)(wbase + 3 * m)).x;
-        qw.s4 = read_imageui(src0_q_img, (int)(wbase + 4 * m)).x;
-        qw.s5 = read_imageui(src0_q_img, (int)(wbase + 5 * m)).x;
-        qw.s6 = read_imageui(src0_q_img, (int)(wbase + 6 * m)).x;
-        qw.s7 = read_imageui(src0_q_img, (int)(wbase + 7 * m)).x;
-
-        for (uint idx = lid; idx < TILESIZE_N * 8; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            const uint c = col_base + t;
-            sh_qa[t][u] = (c < (uint)n_no_padding) ? src1_qa[c * k_u + (step >> 2) + u] : 0u;
-        }
-        if (lid < TILESIZE_N) {
-            const uint c = col_base + lid;
-            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
-        }
-        barrier(CLK_LOCAL_MEM_FENCE);
-
-#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
-        #pragma unroll
-        for (int g = 0; g < NGROUPS; ++g) {
-            const int b = g * 4;
-            float4 rf;
-            rf.s0 = (float)dot8_q8a(qw, sh_qa[b+0]);  rf.s1 = (float)dot8_q8a(qw, sh_qa[b+1]);
-            rf.s2 = (float)dot8_q8a(qw, sh_qa[b+2]);  rf.s3 = (float)dot8_q8a(qw, sh_qa[b+3]);
-            acc[g] += d_w * LD4(sh_d, b) * rf;
-        }
-#undef LD4
-        barrier(CLK_LOCAL_MEM_FENCE);
-    }
-
-    if (!row_valid) {
-        return;
-    }
-
-    #pragma unroll
-    for (int g = 0; g < NGROUPS; ++g) {
-        const uint b = (uint)(g * 4);
-        const float4 a = acc[g];
-        const uint c0 = col_base + b;
-        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
-        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
-        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
-        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
-    }
-#undef NGROUPS
-}

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

@@ -163,95 +163,3 @@ __kernel void kernel_gemv_moe_mxfp4_f32_ns(
     }
 
 }
-
-__attribute__((qcom_reqd_sub_group_size("half")))
-__kernel void kernel_gemv_moe_mxfp4_f32_ns_wimg(
-    __read_only image1d_buffer_t src0_q,
-    __global uchar * src0_e,
-    __read_only image1d_buffer_t src1,
-    __global uint * src2,
-    __global float * dst,
-    ulong         offsetd,
-    int           ne00,
-    int           ne01,
-    int           ne11
-) {
-    uint i01  = get_global_id(0);
-    uint i20  = get_global_id(2);
-    uint sgid = get_local_id(1);
-    uint slid = get_sub_group_local_id();
-
-    if (i01 >= ne01) {
-        return;
-    }
-
-    uint i11 = i20 % ne11;
-
-    uint expert_id = src2[i20];
-    uint expert_offset = expert_id * ne00 * ne01 / 32;
-
-    __private float sum = 0.0f;
-
-    for (uint ib00 = sgid; ib00 < (ne00 / QK_MXFP4); ib00 += N_SIMDGROUP) {
-
-        uint4 regQ;
-        uint block_offset = expert_offset * 4 + ib00 * ne01 * 4 + i01;
-
-        regQ.s0 = read_imageui(src0_q, (int)(block_offset)).x;
-        regQ.s1 = read_imageui(src0_q, (int)(block_offset + ne01)).x;
-        regQ.s2 = read_imageui(src0_q, (int)(block_offset + ne01 * 2)).x;
-        regQ.s3 = read_imageui(src0_q, (int)(block_offset + ne01 * 3)).x;
-
-        uint offset = i11 * ne00 / 4 + ib00 * 8;
-
-        half8 fp16x8 = mxfp4_to_fp16_packed8(as_ushort2(regQ.s0));
-
-        float4 shared_y4;
-        shared_y4 = read_imagef(src1, (offset + 0));
-        float4 acc = shared_y4 * convert_float4(fp16x8.lo);
-
-        shared_y4 = read_imagef(src1, (offset + 1));
-        acc += shared_y4 * convert_float4(fp16x8.hi);
-
-        fp16x8 = mxfp4_to_fp16_packed8(as_ushort2(regQ.s1));
-
-        shared_y4 = read_imagef(src1, (offset + 2));
-        acc += shared_y4 * convert_float4(fp16x8.lo);
-
-        shared_y4 = read_imagef(src1, (offset + 3));
-        acc += shared_y4 * convert_float4(fp16x8.hi);
-
-        fp16x8 = mxfp4_to_fp16_packed8(as_ushort2(regQ.s2));
-
-        shared_y4 = read_imagef(src1, (offset + 4));
-        acc += shared_y4 * convert_float4(fp16x8.lo);
-
-        shared_y4 = read_imagef(src1, (offset + 5));
-        acc += shared_y4 * convert_float4(fp16x8.hi);
-
-        fp16x8 = mxfp4_to_fp16_packed8(as_ushort2(regQ.s3));
-
-        shared_y4 = read_imagef(src1, (offset + 6));
-        acc += shared_y4 * convert_float4(fp16x8.lo);
-
-        shared_y4 = read_imagef(src1, (offset + 7));
-        acc += shared_y4 * convert_float4(fp16x8.hi);
-
-        uchar regE = src0_e[ib00 * ne01 + i01 + expert_offset];
-        sum += e8m0_to_fp32(regE) * ((acc.s0 + acc.s1) + (acc.s2 + acc.s3));
-    }
-
-    __local float reduceLM[SIMDGROUP_WIDTH * (N_SIMDGROUP - 1)];
-    if (sgid == 1) reduceLM[SIMDGROUP_WIDTH * 0 + slid] = sum;
-    if (sgid == 2) reduceLM[SIMDGROUP_WIDTH * 1 + slid] = sum;
-    if (sgid == 3) reduceLM[SIMDGROUP_WIDTH * 2 + slid] = sum;
-    barrier(CLK_LOCAL_MEM_FENCE);
-    if (sgid == 0) sum += reduceLM[SIMDGROUP_WIDTH * 0 + slid];
-    if (sgid == 0) sum += reduceLM[SIMDGROUP_WIDTH * 1 + slid];
-    if (sgid == 0) sum += reduceLM[SIMDGROUP_WIDTH * 2 + slid];
-
-    if (sgid == 0) {
-        dst = dst + (offsetd >> 2);
-        dst[i01 + i20 * ne01] = sum;
-    }
-}

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

@@ -153,114 +153,3 @@ __kernel void kernel_gemv_moe_q4_k_f32_ns(
         dst[i01 + i20 * ne01] = sum;
     }
 }
-
-__attribute__((qcom_reqd_sub_group_size("half")))
-__kernel void kernel_gemv_moe_q4_k_f32_ns_wimg(
-    __read_only image1d_buffer_t src0_q,
-    __global half *         src0_d,
-    __global half *         src0_dm,
-    __global uchar *        src0_s,
-    __read_only image1d_buffer_t src1,
-    __global uint *         src2,
-    __global float *        dst,
-    ulong                   offsetd,
-    int                     ne00,
-    int                     ne01,
-    int                     ne11
-) {
-    uint i01  = get_global_id(0);
-    uint i20  = get_global_id(2);
-    uint sgid = get_local_id(1);
-    uint slid = get_sub_group_local_id();
-
-    if (i01 >= ne01) {
-        return;
-    }
-
-    uint i11 = i20 % ne11;
-
-    uint expert_id = src2[i20];
-
-    int num_superblocks = ne00 / QK_K;
-    int num_subblocks = ne00 / 32;
-    int scales_per_row = num_superblocks * K_SCALE_SIZE;
-
-    uint expert_q_offset = expert_id * (ne00 / 8) * ne01;
-    uint expert_d_offset = expert_id * num_superblocks * ne01;
-
-    __private float sum = 0.0f;
-
-    for (uint ib = sgid; ib < num_subblocks; ib += N_SIMDGROUP) {
-        uint sb = ib / 8;
-        uint j  = ib % 8;
-
-        half d_val   = src0_d[expert_d_offset + sb * ne01 + i01];
-        half dm_val  = src0_dm[expert_d_offset + sb * ne01 + i01];
-
-        global const uchar * sc = src0_s + (expert_id * ne01 + i01) * scales_per_row + sb * K_SCALE_SIZE;
-        uchar sv, mn;
-        get_scale_min_k4(j, sc, &sv, &mn);
-
-        float scale = (float)d_val * (float)sv;
-        float minv  = (float)dm_val * (float)mn;
-
-        uint q_base = expert_q_offset + ib * ne01 * 4 + i01;
-
-        uint4 regQ;
-        regQ.s0 = read_imageui(src0_q, (int)(q_base)).x;
-        regQ.s1 = read_imageui(src0_q, (int)(q_base + ne01)).x;
-        regQ.s2 = read_imageui(src0_q, (int)(q_base + ne01 * 2)).x;
-        regQ.s3 = read_imageui(src0_q, (int)(q_base + ne01 * 3)).x;
-
-        uint y_offset = i11 * ne00 / 4 + ib * 8;
-
-        float8 fp32x8 = q4_k_to_fp32_packed8(as_ushort2(regQ.s0), scale, minv);
-
-        float4 shared_y4;
-        shared_y4 = read_imagef(src1, (y_offset + 0));
-        float4 acc = shared_y4 * fp32x8.lo;
-
-        shared_y4 = read_imagef(src1, (y_offset + 1));
-        acc += shared_y4 * fp32x8.hi;
-
-        fp32x8 = q4_k_to_fp32_packed8(as_ushort2(regQ.s1), scale, minv);
-
-        shared_y4 = read_imagef(src1, (y_offset + 2));
-        acc += shared_y4 * fp32x8.lo;
-
-        shared_y4 = read_imagef(src1, (y_offset + 3));
-        acc += shared_y4 * fp32x8.hi;
-
-        fp32x8 = q4_k_to_fp32_packed8(as_ushort2(regQ.s2), scale, minv);
-
-        shared_y4 = read_imagef(src1, (y_offset + 4));
-        acc += shared_y4 * fp32x8.lo;
-
-        shared_y4 = read_imagef(src1, (y_offset + 5));
-        acc += shared_y4 * fp32x8.hi;
-
-        fp32x8 = q4_k_to_fp32_packed8(as_ushort2(regQ.s3), scale, minv);
-
-        shared_y4 = read_imagef(src1, (y_offset + 6));
-        acc += shared_y4 * fp32x8.lo;
-
-        shared_y4 = read_imagef(src1, (y_offset + 7));
-        acc += shared_y4 * fp32x8.hi;
-
-        sum += ((acc.s0 + acc.s1) + (acc.s2 + acc.s3));
-    }
-
-    __local float reduceLM[SIMDGROUP_WIDTH * (N_SIMDGROUP - 1)];
-    if (sgid == 1) reduceLM[SIMDGROUP_WIDTH * 0 + slid] = sum;
-    if (sgid == 2) reduceLM[SIMDGROUP_WIDTH * 1 + slid] = sum;
-    if (sgid == 3) reduceLM[SIMDGROUP_WIDTH * 2 + slid] = sum;
-    barrier(CLK_LOCAL_MEM_FENCE);
-    if (sgid == 0) sum += reduceLM[SIMDGROUP_WIDTH * 0 + slid];
-    if (sgid == 0) sum += reduceLM[SIMDGROUP_WIDTH * 1 + slid];
-    if (sgid == 0) sum += reduceLM[SIMDGROUP_WIDTH * 2 + slid];
-
-    if (sgid == 0) {
-        dst = dst + (offsetd >> 2);
-        dst[i01 + i20 * ne01] = sum;
-    }
-}

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

@@ -1,36 +0,0 @@
-// Fused MoE combine epilogue: replaces the router-weight MUL + the (n_expert_used-1)
-// cross-expert ADD chain with ONE weighted-sum-across-experts pass.
-//   dst[row, tok] = sum_e experts[row, e, tok] * weights[0, e, tok]
-// experts: [n_embd, n_expert_used, n_tokens] f32 (contiguous after down-proj GEMM)
-// weights: [1, n_expert_used, n_tokens] f32
-// dst:     [n_embd, n_tokens] f32
-// One read of experts + one write of dst (eliminates the intermediate weighted
-// buffer and the k-1 elementwise add round-trips). Vectorized float4 over rows.
-// strides e1/e2/w1/w2/d1 are in ELEMENTS (floats).
-
-__kernel void kernel_moe_combine_f32(
-        __global const char * e_buf, ulong off_e,
-        __global const char * w_buf, ulong off_w,
-        __global       char * d_buf, ulong off_d,
-        int  n_embd4,            // n_embd / 4
-        int  k,                  // n_expert_used
-        int  n_tokens,
-        uint e1, uint e2,        // experts strides (elements): per-expert, per-token
-        uint w1, uint w2,        // weights strides (elements)
-        uint d1)                 // dst per-token stride (elements)
-{
-    const uint r4  = get_global_id(0);
-    const uint tok = get_global_id(1);
-    if (r4 >= (uint)n_embd4 || tok >= (uint)n_tokens) return;
-
-    __global const float * E = (__global const float *)(e_buf + off_e) + tok*e2 + r4*4u;
-    __global const float * W = (__global const float *)(w_buf + off_w) + tok*w2;
-
-    float4 acc = (float4)(0.0f);
-    for (int e = 0; e < k; ++e) {
-        acc = mad(vload4(0, E + (uint)e*e1), (float4)(W[(uint)e*w1]), acc);
-    }
-
-    __global float * D = (__global float *)(d_buf + off_d) + tok*d1 + r4*4u;
-    vstore4(acc, 0, D);
-}

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

@@ -1,64 +0,0 @@
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-
-// Fused MoE activation reorder + q8_1 quantization for the dp4a prefill GEMM.
-// Combines kernel_moe_reorder_b (gather src1 rows per the post-router map) with
-// the q8_1 quant pre-pass, so the f32 reordered-activation tile buffer is never
-// materialised (saves a full write + read of [tok_slots * ne00] floats).
-//
-// One work-item per (token_slot, 32-block). Padding lanes (router 0xFFFFFFFF)
-// emit d=0,s=0,qs=0 so they contribute nothing to the GEMM, exactly as the
-// reorder zero-fill did. Output layout matches kernel_moe_quant_a_q8_1:
-//   qa[token_slot*K + blk*32 + i], da/sa[token_slot*(K/32) + blk].
-__kernel void kernel_moe_reorder_quant_a_q8_1(
-        __global const float  * src,        // original activations (offset applied)
-        __global const uint   * router,     // post-router indices [tok_slots]
-        __global       char   * qa,
-        __global       half   * da,
-        __global       half   * sa,
-        __global const int    * total_tiles,
-        uint  K,
-        ushort map_ratio,
-        uint  tile_size,
-        uint  n_kblocks                      // K / 32
-) {
-    const uint blk = get_global_id(0);       // 32-block along K
-    const uint tok = get_global_id(1);       // token slot (post_router_idx)
-
-    if (blk >= n_kblocks || tok >= (uint)total_tiles[0] * tile_size) {
-        return;
-    }
-
-    const uint out_base = tok * K + blk * 32;
-    const uint bidx     = tok * n_kblocks + blk;
-
-    const uint router_idx = router[tok];
-
-    float v[32];
-    float amax = 0.0f;
-    if (router_idx == 0xFFFFFFFF) {
-        #pragma unroll
-        for (int i = 0; i < 32; ++i) v[i] = 0.0f;
-    } else {
-        const uint act_idx = router_idx / map_ratio;
-        const uint in_base = act_idx * K + blk * 32;
-        #pragma unroll
-        for (int i = 0; i < 32; ++i) {
-            v[i] = src[in_base + i];
-            amax = fmax(amax, fabs(v[i]));
-        }
-    }
-
-    const float d  = amax / 127.0f;
-    const float id = (amax > 0.0f) ? (127.0f / amax) : 0.0f;
-
-    int sum = 0;
-    #pragma unroll
-    for (int i = 0; i < 32; ++i) {
-        const int q = (int)rint(v[i] * id);
-        qa[out_base + i] = (char)q;
-        sum += q;
-    }
-
-    da[bidx] = (half)d;
-    sa[bidx] = (half)(d * (float)sum);
-}

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

@@ -1,42 +0,0 @@
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-
-// Quantize a contiguous [N, K] f32 activation buffer (token-major, K contiguous
-// per token) into q8_1 blocks of 32: int8 quants + per-block scale d + per-block
-// sum s (= d * Sum(qs)). Consumed by kernel_gemm_noshuffle_q4_k_q8_1_dp4a for the
-// dp4a (int8) dense q4_K prefill GEMM. One work-item per 32-element block.
-__kernel void kernel_quant_a_q8_1(
-        __global const float * src,   // [N * K]
-        __global       char  * qa,    // [N * K]
-        __global       half  * da,    // [N * (K/32)]
-        __global       half  * sa,    // [N * (K/32)]
-        int total_blocks              // N * (K/32)
-) {
-    const int blk = get_global_id(0);
-    if (blk >= total_blocks) {
-        return;
-    }
-
-    const int base = blk * 32;
-
-    float v[32];
-    float amax = 0.0f;
-    #pragma unroll
-    for (int i = 0; i < 32; ++i) {
-        v[i] = src[base + i];
-        amax = fmax(amax, fabs(v[i]));
-    }
-
-    const float d  = amax / 127.0f;
-    const float id = (amax > 0.0f) ? (127.0f / amax) : 0.0f;
-
-    int sum = 0;
-    #pragma unroll
-    for (int i = 0; i < 32; ++i) {
-        const int q = (int)rint(v[i] * id);
-        qa[base + i] = (char)q;
-        sum += q;
-    }
-
-    da[blk] = (half)d;
-    sa[blk] = (half)(d * (float)sum);
-}

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

@@ -6501,14 +6501,6 @@ static vk_device ggml_vk_get_device(size_t idx) {
                 device->mul_mat_id_m[i] = true;
                 device->mul_mat_id_s[i] = false;
                 break;
-            case VK_VENDOR_ID_QUALCOMM:
-                device->mul_mat_l[i] = false;
-                device->mul_mat_m[i] = true;
-                device->mul_mat_s[i] = true;
-                device->mul_mat_id_l[i] = false;
-                device->mul_mat_id_m[i] = true;
-                device->mul_mat_id_s[i] = true;
-                break;
 #endif
             default:
                 device->mul_mat_l[i] = true;

ggml/src/ggml.c

@@ -1079,7 +1079,6 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
     "RWKV_WKV7",
     "SOLVE_TRI",
     "GATED_DELTA_NET",
-    "LIGHTNING_INDEXER",
 
     "UNARY",
 
@@ -1097,7 +1096,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
     "GLU",
 };
 
-static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT != 98");
+static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97");
 
 static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "none",
@@ -1191,7 +1190,6 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "rwkv_wkv7(r, w, k, v, a, b, s)",
     "A X = B, A triangular, solve X",
     "gated_delta_net(q, k, v, g, beta, s)",
-    "lightning_indexer(q, k, weights, mask)",
 
     "unary(x)",
 
@@ -1209,7 +1207,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "glu(x)",
 };
 
-static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT != 98");
+static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97");
 
 static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");
 
@@ -6289,42 +6287,6 @@ struct ggml_tensor * ggml_gated_delta_net(
     return result;
 }
 
-// ggml_lightning_indexer
-
-struct ggml_tensor * ggml_lightning_indexer(
-        struct ggml_context * ctx,
-        struct ggml_tensor  * q,
-        struct ggml_tensor  * k,
-        struct ggml_tensor  * weights,
-        struct ggml_tensor  * mask) {
-
-    GGML_ASSERT(       q->type == GGML_TYPE_F32);
-    GGML_ASSERT( weights->type == GGML_TYPE_F32);
-    GGML_ASSERT(    mask->type == GGML_TYPE_F16);
-    GGML_ASSERT(      q->ne[0] == k->ne[0]);
-    GGML_ASSERT(   mask->ne[0] == k->ne[2]);
-    GGML_ASSERT(      q->ne[1] == weights->ne[0]);
-    GGML_ASSERT(      k->ne[1] == 1);
-    GGML_ASSERT(   mask->ne[1] == q->ne[2]);
-    GGML_ASSERT(      q->ne[2] == weights->ne[1]);
-    GGML_ASSERT(weights->ne[2] == 1);
-    GGML_ASSERT(   mask->ne[2] == 1);
-    GGML_ASSERT(      q->ne[3] == k->ne[3]);
-    GGML_ASSERT(      k->ne[3] == weights->ne[3]);
-    GGML_ASSERT(weights->ne[3] % mask->ne[3] == 0);
-
-    int64_t ne[4] = { k->ne[2], q->ne[2], 1, q->ne[3] };
-    struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne);
-
-    result->op   = GGML_OP_LIGHTNING_INDEXER;
-    result->src[0] = q;
-    result->src[1] = k;
-    result->src[2] = weights;
-    result->src[3] = mask;
-
-    return result;
-}
-
 ////////////////////////////////////////////////////////////////////////////////
 
 struct ggml_hash_set ggml_hash_set_new(size_t size) {

ggml/src/gguf.cpp

@@ -557,10 +557,6 @@ static struct gguf_context * gguf_init_from_reader(const struct gguf_reader & gr
                 GGML_LOG_ERROR("%s: encountered bad_alloc error while reading key %" PRIi64 "\n", __func__, i);
                 ok = false;
             }
-            if (ok && key.empty()) {
-                GGML_LOG_ERROR("%s: key %" PRIi64 " is empty\n", __func__, i);
-                ok = false;
-            }
             for (size_t j = 0; ok && j < ctx->kv.size(); ++j) {
                 if (key == ctx->kv[j].key) {
                     GGML_LOG_ERROR("%s: duplicate key '%s' for tensors %zu and %" PRIi64 " \n", __func__, key.c_str(), j, i);

src/llama-context.cpp

@@ -55,12 +55,6 @@ static const llm_fused_op_probe llm_fused_op_gdn_ch_probe = {
     /*.n_tokens_per_seq =*/ 16,
 };
 
-static const llm_fused_op_probe llm_fused_op_lid_probe = {
-    /*.op               =*/ LLM_FUSED_OP_LIGHTNING_INDEXER,
-    /*.name             =*/ "Lightning Indexer",
-    /*.n_tokens_per_seq =*/ 1,
-};
-
 llama_context::llama_context(
         const llama_model & model,
               llama_context_params params) :
@@ -232,9 +226,6 @@ llama_context::llama_context(
     cparams.fused_gdn_ch = true;
     cparams.auto_fgdn    = true;
 
-    cparams.fused_lid    = true;
-    cparams.auto_flid    = true;
-
     // with causal attention, the batch size is limited by the context size
     cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch;
 
@@ -531,12 +522,6 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3
         resolve(llm_fused_op_gdn_ch_probe, cparams.fused_gdn_ch);
         cparams.auto_fgdn = false;
     }
-
-    if (cparams.auto_flid) {
-        LLAMA_LOG_INFO("%s: resolving fused Lightning Indexer support:\n", func);
-        resolve(llm_fused_op_lid_probe, cparams.fused_lid);
-        cparams.auto_flid = false;
-    }
 }
 
 void llama_context::sched_reserve() {

src/llama-cparams.h

@@ -41,8 +41,6 @@ struct llama_cparams {
     bool fused_gdn_ar;       // use fused gated delta net (autoregressive)
     bool fused_gdn_ch;       // use fused gated delta net (chunked)
     bool auto_fgdn;
-    bool fused_lid;          // use fused lightning indexer
-    bool auto_flid;
     bool no_perf;
     bool warmup;             // TODO: remove [TAG_LLAMA_GRAPH_NO_WARMUP]
     bool op_offload;

src/llama-graph.cpp

@@ -842,7 +842,7 @@ static void dsv4_build_comp_inputs(
         GGML_ASSERT(n_stream > 0);
         GGML_ASSERT(n_tokens%n_stream == 0);
 
-        inp.kq_mask = ggml_new_tensor_4d(ctx, (strcmp(name, "lid") != 0 && cparams.flash_attn) || (strcmp(name, "lid") == 0 && cparams.fused_lid) ? GGML_TYPE_F16 : GGML_TYPE_F32, plan.n_kv, n_tokens/n_stream, 1, n_stream);
+        inp.kq_mask = ggml_new_tensor_4d(ctx, cparams.flash_attn && strcmp(name, "lid") != 0 ? GGML_TYPE_F16 : GGML_TYPE_F32, plan.n_kv, n_tokens/n_stream, 1, n_stream);
         ggml_set_input(inp.kq_mask);
         ggml_set_name(inp.kq_mask, (std::string("dsv4_") + name + "_kq_mask").c_str());
     }
@@ -3025,9 +3025,9 @@ llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const {
     {
         inp->self_k_idxs_lid = mctx_cur->get_lid()->build_input_k_idxs(ctx0, ubatch);
 
-        // ensure that mask type matches fused lightning indexer use (requires f16 mask)
+        // ensure F32 mask
         auto cparams_copy = cparams;
-        cparams_copy.flash_attn = cparams.fused_lid;
+        cparams_copy.flash_attn = false;
 
         inp->self_kq_mask_lid = build_attn_inp_kq_mask(ctx0, mctx_cur->get_lid(), ubatch, cparams_copy);
         inp->self_kq_mask_lid_cnv = inp->self_kq_mask_lid;

src/llama-graph.h

@@ -42,7 +42,6 @@ enum llm_fused_op {
     LLM_FUSED_OP_FLASH_ATTN,
     LLM_FUSED_OP_GDN_AR,
     LLM_FUSED_OP_GDN_CH,
-    LLM_FUSED_OP_LIGHTNING_INDEXER,
 };
 
 enum llm_ffn_op_type : int {

src/llama-kv-cache-dsv4.cpp

@@ -29,15 +29,6 @@ static uint32_t dsv4_comp_size(uint32_t kv_size, uint32_t ratio) {
     return std::max<uint32_t>(1, (kv_size + ratio - 1)/ratio);
 }
 
-static void dsv4_clear_tensor_stream(ggml_tensor * tensor, uint32_t stream) {
-    GGML_ASSERT(ggml_is_contiguous(tensor));
-    GGML_ASSERT(tensor->ne[3] == 1);
-    GGML_ASSERT(stream < (uint32_t) tensor->ne[2]);
-
-    const size_t stream_size = tensor->nb[2];
-    ggml_backend_tensor_memset(tensor, 0, stream*stream_size, stream_size);
-}
-
 static int64_t dsv4_stream_offset(uint32_t n_stream, llama_seq_id seq_id, uint32_t size) {
     if (n_stream <= 1) {
         return 0;
@@ -790,20 +781,11 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
             __func__, name, ratio, state_size, n_embd_state, n_stream, layers.size(), total_size()/1024.0/1024.0);
 }
 
-void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) {
+void llama_dsv4_comp_state::clear(bool data) {
     if (!data) {
         return;
     }
 
-    if (seq_id >= 0) {
-        GGML_ASSERT((uint32_t) seq_id < n_stream);
-        for (const auto & layer : layers) {
-            dsv4_clear_tensor_stream(layer.kv,    (uint32_t) seq_id);
-            dsv4_clear_tensor_stream(layer.score, (uint32_t) seq_id);
-        }
-        return;
-    }
-
     for (auto & [_, buf] : ctxs_bufs) {
         ggml_backend_buffer_clear(buf.get(), 0);
     }
@@ -1052,7 +1034,7 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
     // graph does not necessarily overwrite; uninitialized buffer contents would
     // otherwise leak in (instance-specific garbage) and corrupt recall. Zero all
     // compressed buffers up front so reads of un-written rows are deterministic.
-    clear_compressed(-1, true);
+    clear_compressed(true);
 }
 
 llama_memory_context_ptr llama_kv_cache_dsv4::init_batch(
@@ -1165,7 +1147,7 @@ bool llama_kv_cache_dsv4::get_can_shift() const {
 
 void llama_kv_cache_dsv4::clear(bool data) {
     kv_raw->clear(data);
-    clear_compressed(-1, true); // DSV4 compressed buffers must never expose stale/uninit rows
+    clear_compressed(true); // DSV4 compressed buffers must never expose stale/uninit rows
 }
 
 bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
@@ -1187,7 +1169,7 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
     const bool res = kv_raw->seq_rm(seq_id, p0, p1);
 
     if (res) {
-        clear_compressed(seq_id, true);
+        clear_compressed(true);
     }
 
     return res;
@@ -1195,29 +1177,22 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
 
 void llama_kv_cache_dsv4::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
     kv_raw->seq_cp(seq_id_src, seq_id_dst, p0, p1);
+    clear_compressed(true);
 }
 
 void llama_kv_cache_dsv4::seq_keep(llama_seq_id seq_id) {
-    GGML_ASSERT(seq_id >= 0 && (uint32_t) seq_id < n_seq_max);
-
     kv_raw->seq_keep(seq_id);
-
-    for (llama_seq_id id = 0; id < (llama_seq_id) n_seq_max; ++id) {
-        if (id == seq_id) {
-            continue;
-        }
-
-        kv_raw->seq_rm(id, -1, -1);
-        clear_compressed(id, true);
-    }
+    clear_compressed(true);
 }
 
 void llama_kv_cache_dsv4::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
     kv_raw->seq_add(seq_id, p0, p1, shift);
+    clear_compressed(true);
 }
 
 void llama_kv_cache_dsv4::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
     kv_raw->seq_div(seq_id, p0, p1, d);
+    clear_compressed(true);
 }
 
 llama_pos llama_kv_cache_dsv4::seq_pos_min(llama_seq_id seq_id) const {
@@ -1353,32 +1328,13 @@ llama_dsv4_comp_state * llama_kv_cache_dsv4::get_lid_state() const {
     return lid_state.get();
 }
 
-void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) {
-    if (seq_id < 0) {
-        kv_csa->clear(data);
-        kv_hca->clear(data);
-        kv_lid->clear(data);
-    } else {
-        GGML_ASSERT((uint32_t) seq_id < n_seq_max);
-
-        const auto clear_seq = [seq_id, data](llama_kv_cache * kv) {
-            kv->seq_rm(seq_id, -1, -1);
-
-            if (data) {
-                for (uint32_t il : kv->get_layer_ids()) {
-                    dsv4_clear_tensor_stream(kv->get_k_storage(il), (uint32_t) seq_id);
-                }
-            }
-        };
-
-        clear_seq(kv_csa.get());
-        clear_seq(kv_hca.get());
-        clear_seq(kv_lid.get());
-    }
-
-    csa_state->clear(seq_id, data);
-    hca_state->clear(seq_id, data);
-    lid_state->clear(seq_id, data);
+void llama_kv_cache_dsv4::clear_compressed(bool data) {
+    kv_csa->clear(data);
+    kv_hca->clear(data);
+    kv_lid->clear(data);
+    csa_state->clear(data);
+    hca_state->clear(data);
+    lid_state->clear(data);
 }
 
 //

src/llama-kv-cache-dsv4.h

@@ -21,7 +21,7 @@ public:
             const char    * name,
         const llama_memory_i::layer_filter_cb & filter);
 
-    void clear(llama_seq_id seq_id, bool data);
+    void clear(bool data);
 
     uint32_t get_ratio()    const;
     uint32_t get_state_size() const;
@@ -67,8 +67,6 @@ private:
 // DSV4 uses a normal raw/SWA token cache plus compressed K-only block caches.
 // The compressed caches are storage only; DSV4-specific visibility and block
 // planning are handled by llama_kv_cache_dsv4_context / llm_graph_input_dsv4.
-// FIXME: currently the cache only supports non-unified mode even if unified flag is passed
-// FIXME: we currently conflate token_pos and buffer contents. See https://github.com/ggml-org/llama.cpp/pull/25521#discussion_r3558173819
 
 class llama_kv_cache_dsv4 : public llama_memory_i {
 public:
@@ -148,7 +146,7 @@ private:
     std::unique_ptr<llama_dsv4_comp_state> hca_state;
     std::unique_ptr<llama_dsv4_comp_state> lid_state;
 
-    void clear_compressed(llama_seq_id seq_id, bool data);
+    void clear_compressed(bool data);
 };
 
 // DSV4 raw attention only uses the SWA half of kv_raw. The base half is kept

src/llama-model.cpp

@@ -313,7 +313,8 @@ llama_model * llama_model_create(llm_arch arch, const llama_model_params & param
 
     if (model != nullptr) {
         model->arch = arch;
-        if (params.split_mode == LLAMA_SPLIT_MODE_TENSOR && !llm_arch_supports_sm_tensor(arch)) {
+        auto & devices = model->devices;
+        if (!devices.empty() && devices[0].is_meta && !llm_arch_supports_sm_tensor(arch)) {
             throw std::runtime_error(std::string("LLAMA_SPLIT_MODE_TENSOR not implemented for architecture '") + llm_arch_name(arch) + "'");
         }
     }
@@ -335,38 +336,38 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
     const llama_hparams & hparams = ud->model->hparams;
     const std::string tensor_name = tensor->name;
 
-    static const std::regex pattern_q_weight        ("blk\\.\\d*\\.attn_q.weight");
-    static const std::regex pattern_kv_weight       ("blk\\.\\d*\\.attn_(k|v).weight");
-    static const std::regex pattern_qkv_weight      ("blk\\.\\d*\\.attn_qkv.weight");
-    static const std::regex pattern_q_bias          ("blk\\.\\d*\\.attn_q\\.bias");
-    static const std::regex pattern_kv_bias         ("blk\\.\\d*\\.attn_(k|v)\\.bias");
-    static const std::regex pattern_qkv_bias        ("blk\\.\\d*\\.attn_qkv.bias");
-    static const std::regex pattern_qk_norm         ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight");
-    static const std::regex pattern_kv_cache        ("cache_(k|v)_l\\d*");
-    static const std::regex pattern_attn_sinks      ("blk\\.\\d*\\.attn_sinks.weight");
-    static const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight");
-    static const std::regex pattern_attn_out_bias   ("blk\\.\\d*\\.attn_output.bias");
-    static const std::regex pattern_attn_gate_weight("blk\\.\\d*\\.attn_gate.weight");
+    const std::regex pattern_q_weight        ("blk\\.\\d*\\.attn_q.weight");
+    const std::regex pattern_kv_weight       ("blk\\.\\d*\\.attn_(k|v).weight");
+    const std::regex pattern_qkv_weight      ("blk\\.\\d*\\.attn_qkv.weight");
+    const std::regex pattern_q_bias          ("blk\\.\\d*\\.attn_q\\.bias");
+    const std::regex pattern_kv_bias         ("blk\\.\\d*\\.attn_(k|v)\\.bias");
+    const std::regex pattern_qkv_bias        ("blk\\.\\d*\\.attn_qkv.bias");
+    const std::regex pattern_qk_norm         ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight");
+    const std::regex pattern_kv_cache        ("cache_(k|v)_l\\d*");
+    const std::regex pattern_attn_sinks      ("blk\\.\\d*\\.attn_sinks.weight");
+    const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight");
+    const std::regex pattern_attn_out_bias   ("blk\\.\\d*\\.attn_output.bias");
+    const std::regex pattern_attn_gate_weight("blk\\.\\d*\\.attn_gate.weight");
 
-    static const std::regex pattern_ssm_dt          ("blk\\.\\d*\\.ssm_dt.bias");
-    static const std::regex pattern_ssm_a           ("blk\\.\\d*\\.ssm_a");
-    static const std::regex pattern_ssm_alpha       ("blk\\.\\d*\\.ssm_alpha.weight");
-    static const std::regex pattern_ssm_beta        ("blk\\.\\d*\\.ssm_beta.weight");
-    static const std::regex pattern_ssm_beta_alpha  ("blk\\.\\d*\\.ssm_ba.weight");
-    static const std::regex pattern_r_cache         ("cache_r_l\\d*");
-    static const std::regex pattern_s_cache         ("cache_s_l\\d*");
-    static const std::regex pattern_ssm_conv1d      ("blk\\.\\d*\\.ssm_conv1d.weight");
-    static const std::regex pattern_ssm_out_weight  ("blk\\.\\d*\\.ssm_out.weight");
+    const std::regex pattern_ssm_dt          ("blk\\.\\d*\\.ssm_dt.bias");
+    const std::regex pattern_ssm_a           ("blk\\.\\d*\\.ssm_a");
+    const std::regex pattern_ssm_alpha       ("blk\\.\\d*\\.ssm_alpha.weight");
+    const std::regex pattern_ssm_beta        ("blk\\.\\d*\\.ssm_beta.weight");
+    const std::regex pattern_ssm_beta_alpha  ("blk\\.\\d*\\.ssm_ba.weight");
+    const std::regex pattern_r_cache         ("cache_r_l\\d*");
+    const std::regex pattern_s_cache         ("cache_s_l\\d*");
+    const std::regex pattern_ssm_conv1d      ("blk\\.\\d*\\.ssm_conv1d.weight");
+    const std::regex pattern_ssm_out_weight  ("blk\\.\\d*\\.ssm_out.weight");
 
-    static const std::regex pattern_ffn_up_gate_weight("blk\\.\\d*\\.ffn_(up|gate)(_exps)?.weight");
-    static const std::regex pattern_ffn_up_gate_bias  ("blk\\.\\d*\\.ffn_(up|gate)(_exps)?.bias");
-    static const std::regex pattern_ffn_gate_up_weight("blk\\.\\d*\\.ffn_gate_up(_exps)?.weight");
-    static const std::regex pattern_ffn_down_weight   ("blk\\.\\d*\\.ffn_down(_exps)?.weight");
-    static const std::regex pattern_ffn_down_bias     ("blk\\.\\d*\\.ffn_down.bias");
-    static const std::regex pattern_ffn_down_exps_bias("blk\\.\\d*\\.ffn_down_exps.bias");
+    const std::regex pattern_ffn_up_gate_weight("blk\\.\\d*\\.ffn_(up|gate)(_exps)?.weight");
+    const std::regex pattern_ffn_up_gate_bias  ("blk\\.\\d*\\.ffn_(up|gate)(_exps)?.bias");
+    const std::regex pattern_ffn_gate_up_weight("blk\\.\\d*\\.ffn_gate_up(_exps)?.weight");
+    const std::regex pattern_ffn_down_weight   ("blk\\.\\d*\\.ffn_down(_exps)?.weight");
+    const std::regex pattern_ffn_down_bias     ("blk\\.\\d*\\.ffn_down.bias");
+    const std::regex pattern_ffn_down_exps_bias("blk\\.\\d*\\.ffn_down_exps.bias");
 
-    static const std::regex pattern_output_weight("output\\.weight");
-    static const std::regex pattern_output_bias  ("output\\.bias");
+    const std::regex pattern_output_weight("output\\.weight");
+    const std::regex pattern_output_bias  ("output\\.bias");
 
     struct tensor_config {
         ggml_backend_meta_split_axis axis;

src/models/deepseek32.cpp

@@ -301,50 +301,43 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
                 indexer_q = ggml_view_4d(ctx0, indexer_q, indexer_q->ne[0], indexer_q->ne[1], indexer_q->ne[2]/n_stream, n_stream, indexer_q->nb[1], indexer_q->nb[2], indexer_q->nb[3]/n_stream, 0);
                 indexer_weights = ggml_view_4d(ctx0, indexer_weights, indexer_weights->ne[0], indexer_weights->ne[1]/n_stream, indexer_weights->ne[2], n_stream, indexer_weights->nb[1], indexer_weights->nb[2]/n_stream, indexer_weights->nb[3]/n_stream, 0);
 
+                // calculate indexer kq
+                indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);
+                cb(indexer_q, "indexer_q", il);
+                indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);
+                cb(indexer_k, "indexer_k", il);
+
+                ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);
+                cb(indexer_kq, "indexer_kq", il);
+
+                // ReLU requires contiguous tensors
+                indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));
+                cb(indexer_kq, "indexer_kq", il);
+
+                // apply ReLU
+                ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq);
+                cb(indexer_score, "indexer_score", il);
+
                 // pre-scale weights to avoid scaling operations on huge indexer_score tensor
                 indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f / sqrtf(float(n_embd_indexer_head * n_indexer_head)));
                 cb(indexer_weights, "indexer_weights", il);
 
-                ggml_tensor * indexer_score = nullptr;
-                if (cparams.fused_lid) {
-                    indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn_dsa->get_kq_mask_lid());
-                    cb(indexer_score, "indexer_score", il);
-                    res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il});
-                } else {
-                    // calculate indexer kq
-                    indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);
-                    cb(indexer_q, "indexer_q", il);
-                    indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);
-                    cb(indexer_k, "indexer_k", il);
+                // multiply scores by indexer weights
+                indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);
+                cb(indexer_score, "indexer_score", il);
 
-                    ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);
-                    cb(indexer_kq, "indexer_kq", il);
+                // sum by q n_indexer_head dimension
+                indexer_score = ggml_sum_rows(ctx0, indexer_score);
+                cb(indexer_score, "indexer_score", il);
 
-                    // ReLU requires contiguous tensors
-                    indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));
-                    cb(indexer_kq, "indexer_kq", il);
+                // permute result to match KQ mask
+                indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
+                cb(indexer_score, "indexer_score", il);
 
-                    // apply ReLU
-                    indexer_score = ggml_relu(ctx0, indexer_kq);
-                    cb(indexer_score, "indexer_score", il);
-
-                    // multiply scores by indexer weights
-                    indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);
-                    cb(indexer_score, "indexer_score", il);
-
-                    // sum by q n_indexer_head dimension
-                    indexer_score = ggml_sum_rows(ctx0, indexer_score);
-                    cb(indexer_score, "indexer_score", il);
-
-                    // permute result to match KQ mask
-                    indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
-                    cb(indexer_score, "indexer_score", il);
-
-                    // mask indexer scores
-                    ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid();
-                    indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask);
-                    cb(indexer_score, "indexer_score", il);
-                }
+                // mask indexer scores
+                ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid();
+                indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask);
+                cb(indexer_score, "indexer_score", il);
 
                 // get indices of top k indexer scores
                 uint32_t n_top_k = indexer_score->ne[0] < n_indexer_top_k ? indexer_score->ne[0] : n_indexer_top_k;

src/models/deepseek4.cpp

@@ -556,32 +556,25 @@ ggml_tensor * llama_model_deepseek4::graph::build_lid_top_k(
             indexer_weights->ne[0], indexer_weights->ne[1]/n_stream, indexer_weights->ne[2], n_stream,
             indexer_weights->nb[1], indexer_weights->nb[2]/n_stream, indexer_weights->nb[3]/n_stream, 0);
 
-    ggml_tensor * indexer_score = nullptr;
-    if (cparams.fused_lid) {
-        indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_lid.kq_mask);
-        cb(indexer_score, "lid_score_masked", il);
-        res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il});
-    } else {
-        indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);
-        cb(indexer_q, "lid_q", il);
-        indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);
-        cb(indexer_k, "lid_k", il);
+    indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);
+    cb(indexer_q, "lid_q", il);
+    indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);
+    cb(indexer_k, "lid_k", il);
 
-        ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);
-        cb(indexer_kq, "lid_kq", il);
+    ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);
+    cb(indexer_kq, "lid_kq", il);
 
-        indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));
-        cb(indexer_kq, "lid_kq", il);
+    indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));
+    cb(indexer_kq, "lid_kq", il);
 
-        indexer_score = ggml_relu(ctx0, indexer_kq);
-        indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);
-        indexer_score = ggml_sum_rows(ctx0, indexer_score);
-        indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
-        cb(indexer_score, "lid_score", il);
+    ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq);
+    indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);
+    indexer_score = ggml_sum_rows(ctx0, indexer_score);
+    indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
+    cb(indexer_score, "lid_score", il);
 
-        indexer_score = ggml_add(ctx0, indexer_score, inp_lid.kq_mask);
-        cb(indexer_score, "lid_score_masked", il);
-    }
+    indexer_score = ggml_add(ctx0, indexer_score, inp_lid.kq_mask);
+    cb(indexer_score, "lid_score_masked", il);
 
     const uint32_t n_top_k = indexer_score->ne[0] < hparams.indexer_top_k ? indexer_score->ne[0] : hparams.indexer_top_k;
     ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, indexer_score, n_top_k));

tests/test-backend-ops.cpp

@@ -7097,67 +7097,6 @@ struct test_diag : public test_case {
     }
 };
 
-// GGML_OP_LIGHTNING_INDEXER
-struct test_lightning_indexer : public test_case {
-    const int64_t hsk; // indexer K head size
-    const int64_t nh; // num indexer heads
-    const int64_t kv; // kv size
-    const int64_t nb; // batch size
-    const int64_t ns; // num streams
-    const int64_t nm; // ne[3] of mask
-
-    const ggml_type type_K;
-
-    std::string vars() override {
-        return VARS_TO_STR7(hsk, nh, kv, nb, ns, nm, type_K);
-    }
-
-    double max_nmse_err() override {
-        return 1e-6;
-    }
-
-    uint64_t op_flops(ggml_tensor * t) override {
-        GGML_UNUSED(t);
-        return ((2 * hsk + 2) * nh + 1) * kv * nb * ns;
-    }
-
-    test_lightning_indexer(int64_t hsk = 128, int64_t nh = 64, int64_t kv = 256, int64_t nb = 128, int64_t ns = 1, int64_t nm = 1, ggml_type type_K = GGML_TYPE_F16)
-        : hsk(hsk), nh(nh), kv(kv), nb(nb), ns(ns), nm(nm), type_K(type_K) {}
-
-    ggml_tensor * build_graph(ggml_context * ctx) override {
-        ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, hsk, nh, nb, ns);
-        ggml_set_param(q);
-        ggml_set_name(q, "q");
-
-        ggml_tensor * k = ggml_new_tensor_4d(ctx, type_K, hsk, 1, kv, ns);
-        ggml_set_param(k);
-        ggml_set_name(k, "k");
-
-        ggml_tensor * w = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, nh, nb, 1, ns);
-        ggml_set_param(w);
-        ggml_set_name(w, "w");
-
-        ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kv, nb, 1, nm);
-        ggml_set_param(m);
-        ggml_set_name(m, "m");
-
-        ggml_tensor * out = ggml_lightning_indexer(ctx, q, k, w, m);
-        ggml_set_name(out, "out");
-
-        return out;
-    }
-
-    void initialize_tensors(ggml_context * ctx) override {
-        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
-            if (strcmp(t->name, "m") == 0) {
-                init_tensor_kq_mask(t);
-            } else {
-                init_tensor_uniform(t);
-            }
-        }
-    }
-};
-
 // Deserializable generic test case
 struct input_tensor {
     ggml_type type;
@@ -9454,19 +9393,6 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_falcon(2));
 #endif
 
-    // lightning_indexer
-    for (int kv : { 256 }) {
-        for (int bs : { 1, 512 }) {
-            for (int nh : { 32, 64 }) {
-                for (auto [ns, nm] : { std::pair{1, 1}, std::pair{4, 4}, std::pair{4, 1} }) {
-                    for (ggml_type type_K : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0, GGML_TYPE_IQ4_NL}) {
-                        test_cases.emplace_back(new test_lightning_indexer(128, nh, kv, bs, ns, nm, type_K));
-                    }
-                }
-            }
-        }
-    }
-
     return test_cases;
 }
 #ifdef _MSC_VER
@@ -9796,19 +9722,6 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
     test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 1024, 1)); // 4h PP-1024
     test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 64, 1, 1, false, true)); // KDA PP-64
 
-    // lightning_indexer
-    for (int kv : { 256, 4096, 65536 }) {
-        for (int bs : { 1, 512, 2048 }) {
-            for (int nh : { 32, 64 }) {
-                for (int ns : { 1, 4 }) {
-                    for (ggml_type type_K : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0, GGML_TYPE_IQ4_NL}) {
-                        test_cases.emplace_back(new test_lightning_indexer(128, nh, kv, bs, ns, ns, type_K));
-                    }
-                }
-            }
-        }
-    }
-
     return test_cases;
 }
 

tests/test-gguf.cpp

@@ -26,7 +26,6 @@ enum handcrafted_file_type {
     HANDCRAFTED_HEADER_EMPTY               = 800,
 
     HANDCRAFTED_KV_BAD_KEY_SIZE            =  10 + offset_has_kv,
-    HANDCRAFTED_KV_EMPTY_KEY               =  15 + offset_has_kv,
     HANDCRAFTED_KV_BAD_TYPE                =  20 + offset_has_kv,
     // HANDCRAFTED_KV_BAD_VALUE_SIZE          =  30 + offset_has_kv, // removed because it can result in allocations > 1 TB (default sanitizer limit)
     HANDCRAFTED_KV_DUPLICATE_KEY           =  40 + offset_has_kv,
@@ -65,7 +64,6 @@ static std::string handcrafted_file_type_name(const enum handcrafted_file_type h
         case HANDCRAFTED_HEADER_EMPTY:               return "HEADER_EMPTY";
 
         case HANDCRAFTED_KV_BAD_KEY_SIZE:            return "KV_BAD_KEY_SIZE";
-        case HANDCRAFTED_KV_EMPTY_KEY:               return "KV_EMPTY_KEY";
         case HANDCRAFTED_KV_BAD_TYPE:                return "KV_BAD_TYPE";
         case HANDCRAFTED_KV_DUPLICATE_KEY:           return "KV_DUPLICATE_KEY";
         case HANDCRAFTED_KV_BAD_ALIGN:               return "KV_BAD_ALIGN";
@@ -286,9 +284,7 @@ static FILE * get_handcrafted_file(const unsigned int seed, const enum handcraft
         const enum gguf_type type     = gguf_type(hft == HANDCRAFTED_KV_BAD_TYPE ? GGUF_TYPE_COUNT : kv_types[i].first);
         const enum gguf_type type_arr = gguf_type(hft == HANDCRAFTED_KV_BAD_TYPE ? GGUF_TYPE_COUNT : kv_types[i].second);
 
-        const std::string key = hft == HANDCRAFTED_KV_EMPTY_KEY
-            ? ""
-            : "my_key_" + std::to_string((hft == HANDCRAFTED_KV_DUPLICATE_KEY ? i/2 : i));
+        const std::string key = "my_key_" + std::to_string((hft == HANDCRAFTED_KV_DUPLICATE_KEY ? i/2 : i));
 
         if (hft == HANDCRAFTED_KV_BAD_KEY_SIZE) {
             const uint64_t n = -1;
@@ -736,7 +732,6 @@ static std::pair<int, int> test_handcrafted_file(const unsigned int seed) {
         HANDCRAFTED_HEADER_EMPTY,
 
         HANDCRAFTED_KV_BAD_KEY_SIZE,
-        HANDCRAFTED_KV_EMPTY_KEY,
         HANDCRAFTED_KV_BAD_TYPE,
         HANDCRAFTED_KV_DUPLICATE_KEY,
         HANDCRAFTED_KV_BAD_ALIGN,

tests/test-save-load-state.cpp

@@ -78,84 +78,7 @@ static llama_tokens test_baseline(struct llama_model * model, const struct commo
 }
 
 
-// Test 2: sequence removal isolation
-// - decode the same prefix into two sequences
-// - remove sequence 0
-// - verify that sequence 1 remains unchanged
-static bool test_seq_rm_isolated(
-        struct llama_model         * model,
-        const struct common_params & params,
-        const llama_tokens         & tokens) {
-    auto params_ctx = common_context_params_to_llama(params);
-    params_ctx.n_ctx      = 256;
-    params_ctx.n_seq_max  = 2;
-    params_ctx.kv_unified = true;
-
-    auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
-    if (!ctx) {
-        LOG_ERR("%s: failed to create context\n", __func__);
-        return false;
-    }
-
-    LOG("\n=== Test 2: sequence removal isolation ===\n");
-
-    const size_t n_tokens = tokens.size() < 128 ? tokens.size() : 128;
-    for (llama_seq_id seq_id = 0; seq_id < 2; ++seq_id) {
-        llama_batch_ptr batch(n_tokens, 0, 1);
-        for (size_t i = 0; i < n_tokens; ++i) {
-            common_batch_add(batch.get(), tokens[i], i, { seq_id }, false);
-        }
-
-        if (llama_decode(ctx.get(), batch.get())) {
-            LOG_ERR("%s: failed to decode prompt for sequence %d\n", __func__, seq_id);
-            return false;
-        }
-    }
-
-    const auto get_seq_state = [&](llama_seq_id seq_id, std::vector<uint8_t> & state) {
-        const size_t state_size = llama_state_seq_get_size(ctx.get(), seq_id);
-        if (state_size == 0) {
-            LOG_ERR("%s: sequence state is empty\n", __func__);
-            return false;
-        }
-
-        state.resize(state_size);
-        const size_t ncopy = llama_state_seq_get_data(ctx.get(), state.data(), state.size(), seq_id);
-        if (ncopy != state.size()) {
-            LOG_ERR("%s: sequence state length %zu does not match expected length %zu\n",
-                    __func__, ncopy, state.size());
-            return false;
-        }
-
-        return true;
-    };
-
-    std::vector<uint8_t> state_before;
-    if (!get_seq_state(1, state_before)) {
-        return false;
-    }
-
-    if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) {
-        LOG_ERR("%s: failed to remove sequence 0\n", __func__);
-        return false;
-    }
-
-    std::vector<uint8_t> state_after;
-    if (!get_seq_state(1, state_after)) {
-        return false;
-    }
-
-    if (state_before != state_after) {
-        LOG_ERR("%s: removing sequence 0 changed sequence 1\n", __func__);
-        return false;
-    }
-
-    LOG("PASS\n");
-    return true;
-}
-
-
-// Test 3: state load
+// Test 2: state load
 // - create a new context
 // - load state from file
 // - replay the last prompt token
@@ -167,7 +90,7 @@ static bool test_state_load(struct llama_model * model, const struct common_para
     auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
     llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
 
-    LOG("\n=== Test 3: state load ===\n");
+    LOG("\n=== Test 2: state load ===\n");
 
     // Load state from file
     llama_tokens unused_sts(tokens.size());
@@ -203,7 +126,7 @@ static bool test_state_load(struct llama_model * model, const struct common_para
 }
 
 
-// Test 4: seq copy (host)
+// Test 3: seq copy (host)
 // - create a multi-seq context
 // - load state from file
 // - replay the last prompt token
@@ -218,7 +141,7 @@ static bool test_seq_cp_host(struct llama_model * model, const struct common_par
     auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
     llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
 
-    LOG("\n=== Test 4: seq copy (host) ===\n");
+    LOG("\n=== Test 3: seq copy (host) ===\n");
 
     // Load state from file
     llama_tokens unused_sts(tokens.size());
@@ -275,7 +198,7 @@ static bool test_seq_cp_host(struct llama_model * model, const struct common_par
 }
 
 
-// Test 5: seq copy (device)
+// Test 4: seq copy (device)
 // - create a multi-seq context
 // - load state from file
 // - replay the last prompt token
@@ -290,7 +213,7 @@ static bool test_seq_cp_device(struct llama_model * model, const struct common_p
     auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
     llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
 
-    LOG("\n=== Test 5: seq copy (device) ===\n");
+    LOG("\n=== Test 4: seq copy (device) ===\n");
 
     // Load state from file
     llama_tokens unused_sts(tokens.size());
@@ -414,22 +337,17 @@ int main(int argc, char ** argv) {
         return 1;
     }
 
-    // Test 2: sequence removal isolation
-    if (!test_seq_rm_isolated(model, params, tokens)) {
-        return 1;
-    }
-
-    // Test 3: state load
+    // Test 2: state load
     if (!test_state_load(model, params, tokens, result_baseline)) {
         return 1;
     }
 
-    // Test 4: seq copy (host)
+    // Test 3: seq copy (host)
     if (!test_seq_cp_host(model, params, tokens, result_baseline)) {
         return 1;
     }
 
-    // Test 5: seq copy (device)
+    // Test 4: seq copy (device)
     if (!test_seq_cp_device(model, params, tokens, result_baseline)) {
         return 1;
     }

tools/mtmd/clip-graph.h

@@ -20,8 +20,8 @@ struct clip_graph {
     const clip_hparams & hparams;
     projector_type proj_type;
 
-    const clip_image_f32 & img; // for backward compat
-    const clip_image_f32_batch * img_batch = nullptr;
+    // we only support single image per batch
+    const clip_image_f32 & img;
 
     const int patch_size;
     const int n_patches_x;
@@ -63,12 +63,6 @@ struct clip_graph {
     //
     void cb(ggml_tensor * cur0, const char * name, int il) const;
 
-    const clip_image_f32 & get_img(size_t idx) const {
-        GGML_ASSERT(img_batch);
-        GGML_ASSERT(idx < img_batch->entries.size());
-        return img_batch->entries[idx];
-    }
-
     // siglip2 naflex
     ggml_tensor * resize_position_embeddings(uint32_t interpolation_mode = DEFAULT_INTERPOLATION_MODE);
 

tools/mtmd/clip-model.h

@@ -69,7 +69,6 @@ struct clip_hparams {
     std::vector<clip_image_size> image_res_candidates;
     int32_t preproc_min_tiles = 0;
     int32_t preproc_max_tiles = 0;
-    int32_t preproc_tile_size = 0; // local tile size (deepseek-ocr)
     resize_algo image_resize_algo_rf = RESIZE_ALGO_BICUBIC;
     resize_algo image_resize_algo_ov = RESIZE_ALGO_BILINEAR;
     pad_style image_pad_rf = PAD_CEIL;  // padding style for the refined image (e.g. llava-1.6)

tools/mtmd/clip.cpp

@@ -1024,8 +1024,6 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
             GGML_ABORT("missing cgraph builder");
     }
 
-    builder->img_batch = &imgs;
-
     // TODO [QWEN_VIDEO]: improve this in the future
     builder->n_batch = imgs.entries.size();
 
@@ -1582,16 +1580,7 @@ struct clip_model_loader {
                         get_u32(KEY_SAM_N_HEAD, hparams.sam_n_head, true);
                         get_u32(KEY_SAM_N_EMBD, hparams.sam_n_embd, true);
                         get_u32(KEY_ATTN_WINDOW_SIZE, hparams.attn_window_size, true);
-                        hparams.preproc_min_tiles = 2;
-                        if (model.proj_type == PROJECTOR_TYPE_DEEPSEEKOCR) {
-                            hparams.preproc_max_tiles = 9;
-                            hparams.preproc_tile_size = 640;
-                            // the CLIP/ViT body runs its layernorms at 1e-5 (the SAM stage uses 1e-6)
-                            hparams.eps = 1e-5f;
-                        }
                         if (model.proj_type == PROJECTOR_TYPE_DEEPSEEKOCR2) {
-                            hparams.preproc_max_tiles = 6;
-                            hparams.preproc_tile_size = 768;
                             // qwen2 encoder is GQA, requires KEY_N_HEAD_KV
                             get_u32(string_format(KEY_N_HEAD_KV, "vision"), hparams.n_head_kv);
                         }
@@ -3262,9 +3251,6 @@ int clip_n_output_tokens_x(const clip_ctx * ctx, const clip_image_f32 * img) {
             return (img->nx() / params.patch_size) / 2;
         case PROJECTOR_TYPE_STEP3VL:
             return img->nx() / (params.patch_size * params.n_merge);
-        case PROJECTOR_TYPE_DEEPSEEKOCR:
-        case PROJECTOR_TYPE_DEEPSEEKOCR2:
-            return (img->nx() / params.patch_size) / 4;
         default:
             break;
     }
@@ -3474,17 +3460,10 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
             // E.g., 64x64 -> 16x16 patches
             n_patches /= 16;
 
-            if (img->add_viewsep) {
-                // global view: one image-newline per token-row + trailing view separator
-                const int h = static_cast<int>(std::sqrt(static_cast<float>(n_patches)));
-                n_patches = h * (h + 1) + 1;
-            } else if (img->ny() >= img->nx() && img->ny() % img->nx() == 0) {
-                // tile row: one image-newline per token-row
-                const int grid_w = img->ny() / img->nx();
-                const int tile_patches = img->nx() / (patch_size * 4); // patches per tile side (SAM divides by 4)
-                const int h = tile_patches;
-                n_patches = (tile_patches * grid_w + 1) * h;
-            }
+            // build_global_local_features adds image newlines and view separator
+            // Formula: h*(w+1) + 1 where h = w = sqrt(n_patches)
+            int h = static_cast<int>(std::sqrt(static_cast<float>(n_patches)));
+            n_patches = h * (h + 1) + 1;
         } break;
         case PROJECTOR_TYPE_HUNYUANVL:
             {
@@ -4124,10 +4103,7 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32
         case PROJECTOR_TYPE_DEEPSEEKOCR:
         case PROJECTOR_TYPE_DEEPSEEKOCR2:
             {
-                GGML_ASSERT(
-                    (pos_w == pos_h) // overview image
-                    || (pos_h >= pos_w && pos_h % pos_w == 0) // tile images
-                );
+                GGML_ASSERT(pos_w == pos_h);
 
                 const int window = hparams.attn_window_size;
                 const int pos = pos_w;

tools/mtmd/models/deepseekocr.cpp

@@ -96,8 +96,6 @@ ggml_tensor * clip_graph_deepseekocr::build_sam(ggml_tensor * inp_raw) {
     const int n_heads = hparams.sam_n_head;
     const int d_heads = n_embd / n_heads;
     const int window  = hparams.attn_window_size;
-    // SAM stage runs its layernorms at 1e-6
-    const float sam_eps = 1e-6f;
 
     ggml_tensor * inpL;
 
@@ -136,7 +134,7 @@ ggml_tensor * clip_graph_deepseekocr::build_sam(ggml_tensor * inp_raw) {
         ggml_tensor * shortcut = cur;
 
         // layernorm1
-        cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, sam_eps, il);
+        cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il);
 
         const int64_t w0 = cur->ne[1];
         const int64_t h0 = cur->ne[2];
@@ -216,7 +214,7 @@ ggml_tensor * clip_graph_deepseekocr::build_sam(ggml_tensor * inp_raw) {
         ggml_tensor * inpFF = cur;
 
         // layernorm2
-        cur = build_norm(inpFF, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, sam_eps, il);
+        cur = build_norm(inpFF, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il);
 
         // ffn
         cur = build_ffn(cur, layer.ff_up_w, layer.ff_up_b, nullptr, nullptr, layer.ff_down_w, layer.ff_down_b,
@@ -231,12 +229,12 @@ ggml_tensor * clip_graph_deepseekocr::build_sam(ggml_tensor * inp_raw) {
 
     cur = ggml_conv_2d(ctx0, model.neck_0_w, cur, 1, 1, 0, 0, 1, 1);
     cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 1, 2, 0, 3));
-    cur = build_norm(cur, model.neck_1_w, model.neck_1_b, NORM_TYPE_NORMAL, sam_eps, -1);
+    cur = build_norm(cur, model.neck_1_w, model.neck_1_b, NORM_TYPE_NORMAL, hparams.eps, -1);
     cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 2, 0, 1, 3));
 
     cur = ggml_conv_2d(ctx0, model.neck_2_w, cur, 1, 1, 1, 1, 1, 1);
     cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 1, 2, 0, 3));
-    cur = build_norm(cur, model.neck_3_w, model.neck_3_b, NORM_TYPE_NORMAL, sam_eps, -1);
+    cur = build_norm(cur, model.neck_3_w, model.neck_3_b, NORM_TYPE_NORMAL, hparams.eps, -1);
     cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 2, 0, 1, 3));
 
     cur = ggml_conv_2d(ctx0, model.net_2, cur, 2, 2, 1, 1, 1, 1);
@@ -250,40 +248,8 @@ ggml_tensor * clip_graph_deepseekocr::build_sam(ggml_tensor * inp_raw) {
 ggml_cgraph * clip_graph_deepseekocr::build() {
     // patch embedding
     ggml_tensor * inp_raw = build_inp_raw();
-
-    bool is_overview = img.add_viewsep;
-    int n_tiles_per_row = 0;
-
-    // note: we expect either a batch of rows or a batch of overviews, but not a mix of both
-
-    if (!is_overview) {
-        // handle the case where we have a batch of rows
-        // sanity check
-        for (auto & entry : img_batch->entries) {
-            if (entry.add_viewsep) {
-                throw std::runtime_error("DeepSeek-OCR: mixed overview and non-overview images in batch");
-            }
-            if (entry.nx() != img.nx() || entry.ny() != img.ny()) {
-                throw std::runtime_error("DeepSeek-OCR: mixed image sizes in batch");
-            }
-        }
-
-        GGML_ASSERT(img.ny() >= img.nx());
-        GGML_ASSERT(img.ny() % img.nx() == 0);
-        n_tiles_per_row = img.ny() / img.nx();
-
-        // input shape: [tile_size, tile_size * n_tiles_per_row, 3]
-        // we want to reshape it to [tile_size, tile_size, 3, n_tiles_per_row]
-        inp_raw = ggml_reshape_4d(ctx0, inp_raw, img.nx(), img.nx(), n_tiles_per_row, 3);
-        inp_raw = ggml_cont(ctx0, ggml_permute(ctx0, inp_raw, 0, 1, 3, 2));
-    }
-
     ggml_tensor * sam_out = build_sam(inp_raw);
 
-    if (!is_overview) {
-        n_batch = n_tiles_per_row;
-    }
-
     const int clip_n_patches = sam_out->ne[0] * sam_out->ne[1];
 
     ggml_tensor * clip_out;
@@ -291,9 +257,7 @@ ggml_cgraph * clip_graph_deepseekocr::build() {
     {
         ggml_tensor * inp;
 
-        // sam_out: [patch_h, patch_w, n_embd, n_batch]
-        // -> [n_embd, clip_n_patches, n_batch]
-        inp = ggml_reshape_3d(ctx0, sam_out, clip_n_patches, sam_out->ne[2], sam_out->ne[3]);
+        inp = ggml_reshape_2d(ctx0, sam_out, clip_n_patches, sam_out->ne[2]);
         inp = ggml_cont(ctx0, ggml_permute(ctx0, inp, 1, 0, 2, 3));
 
         ggml_tensor * new_pos_embd = model.position_embeddings;
@@ -317,11 +281,8 @@ ggml_cgraph * clip_graph_deepseekocr::build() {
             n_pos        = tgt_size * tgt_size + 1;
         }
 
-        // add CLS token per batch item
-        // inp: [n_embd, clip_n_patches, n_batch]
-        // class_embedding: [n_embd] -> [n_embd, 1, n_batch]
-        ggml_tensor * cls_embd = ggml_repeat_4d(ctx0, model.class_embedding, n_embd, 1, n_batch, 1);
-        inp = ggml_concat(ctx0, cls_embd, inp, 1);
+        // add CLS token
+        inp = ggml_concat(ctx0, model.class_embedding, inp, 1);
 
         // for selecting learned pos embd, used by ViT
         ggml_tensor * positions        = ggml_cast(ctx0, ggml_arange(ctx0, 0, n_pos, 1), GGML_TYPE_I32);
@@ -333,56 +294,25 @@ ggml_cgraph * clip_graph_deepseekocr::build() {
         clip_out = cur;
     }
 
-    // sam_out: [patch_h, patch_w, n_embd, n_batch]
-    // -> [n_embd, clip_n_patches, n_batch]
     sam_out  = ggml_cont(ctx0, ggml_permute(ctx0, sam_out, 1, 2, 0, 3));
-    sam_out  = ggml_reshape_3d(ctx0, sam_out, sam_out->ne[0], clip_n_patches, n_batch);
-
-    // clip_out: [n_embd, n_pos, n_batch] where n_pos = clip_n_patches + 1 (CLS)
-    // strip CLS token: skip first position, view only the patch tokens
-    clip_out = ggml_view_3d(ctx0, clip_out, n_embd, clip_n_patches, n_batch,
-                            clip_out->nb[1], clip_out->nb[2], clip_out->nb[1]);
+    sam_out  = ggml_reshape_2d(ctx0, sam_out, sam_out->ne[0], clip_n_patches);
+    clip_out = ggml_view_2d(ctx0, clip_out, n_embd, clip_n_patches, clip_out->nb[1], clip_out->nb[1]);
 
     ggml_tensor * cur;
     cur = ggml_concat(ctx0, clip_out, sam_out, 0);
     cur = ggml_mul_mat(ctx0, model.mm_fc_w, cur);
     cur = ggml_add(ctx0, cur, model.mm_fc_b);
 
-    if (is_overview) {
-        // global view: weave one newline per row + trailing view separator
-        const auto h     = static_cast<int>(std::sqrt(static_cast<float>(cur->ne[1])));
-        const auto w     = h;
-        const auto n_dim = cur->ne[0];
+    const auto h     = static_cast<int>(std::sqrt(static_cast<float>(cur->ne[1])));
+    const auto w     = h;
+    const auto n_dim = cur->ne[0];
 
-        ggml_tensor * imgnl = ggml_repeat_4d(ctx0, model.image_newline, n_dim, 1, h, 1);
-        cur = ggml_reshape_3d(ctx0, cur, n_dim, w, h);
-        cur = ggml_reshape_2d(ctx0, ggml_concat(ctx0, cur, imgnl, 1), n_dim, (w + 1) * h);
-        cur = ggml_concat(ctx0, cur, model.view_seperator, 1);  // (n_dim, h*(w+1) + 1)
-    } else {
-        // tile row: interleave tiles within each row, add newline per row
-        const int grid_x      = static_cast<int>(std::sqrt(static_cast<float>(clip_n_patches)));
-        const int grid_y      = grid_x;
-        const auto n_dim      = cur->ne[0];
+    ggml_tensor * imgnl;
 
-        // (n_dim, clip_n_patches, n_batch) -> (n_dim, grid_x, grid_y, n_batch)
-        cur = ggml_reshape_4d(ctx0, cur, n_dim, grid_x, grid_y, n_batch);
-
-        // tiles: re-order from A.row0 A.row1 B.row0 B.row1 ...
-        //        to A.row0 B.row0 A.row1 B.row1 ...
-        //        then add nl: A.row0 B.row0 [nl] A.row1 B.row1 [nl] ...
-        // interleave tiles: (n_dim, grid_x, grid_y, n_batch) -> (n_dim, grid_x, n_batch, grid_y)
-        cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 0, 1, 3, 2));
-
-        // merge: (n_dim, grid_x, n_batch, grid_y) -> (n_dim, grid_x*n_batch, grid_y, 1)
-        cur = ggml_reshape_4d(ctx0, cur, n_dim, grid_x * n_batch, grid_y, 1);
-
-        // append newline per row: (n_dim, grid_x*n_batch+1, grid_y, 1)
-        ggml_tensor * imgnl = ggml_repeat_4d(ctx0, model.image_newline, n_dim, 1, grid_y, 1);
-        cur = ggml_concat(ctx0, cur, imgnl, 1);
-
-        // flatten: (n_dim, (grid_x*n_batch+1)*grid_y)
-        cur = ggml_reshape_2d(ctx0, cur, n_dim, (grid_x * n_batch + 1) * grid_y);
-    }
+    imgnl = ggml_repeat_4d(ctx0, model.image_newline, n_dim, 1, h, 1);
+    cur   = ggml_reshape_3d(ctx0, cur, n_dim, w, h);
+    cur   = ggml_reshape_2d(ctx0, ggml_concat(ctx0, cur, imgnl, 1), n_dim, (w + 1) * h);
+    cur   = ggml_concat(ctx0, cur, model.view_seperator, 1);  // (n_dim, h*(w+1) + 1)
 
     cb(cur, "dsocr_output", -1);
 

tools/mtmd/models/models.h

@@ -127,7 +127,6 @@ struct clip_graph_deepseekocr : clip_graph {
     clip_graph_deepseekocr(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
     ggml_cgraph * build() override;
     ggml_tensor * build_sam(ggml_tensor * inp); // build the SAM model
-    // bool support_batch() const override { return true; } // TODO: support batch for DeepSeek-OCR v1
 };
 
 struct clip_graph_deepseekocr2 : clip_graph_deepseekocr {

tools/mtmd/mtmd-image.cpp

@@ -1107,7 +1107,44 @@ mtmd_image_preproc_out mtmd_image_preprocessor_internvl::preprocess(const clip_i
 // mtmd_image_preprocessor_deepseekocr
 //
 
-std::vector<clip_image_size> mtmd_image_preprocessor_deepseekocr::get_target_ratios() const {
+mtmd_image_preproc_out mtmd_image_preprocessor_deepseekocr::preprocess(const clip_image_u8 & img) {
+    static constexpr int native_resolutions[] = { 1024 /* base */, 1280 /* large */ };
+    // TODO: support 512 (tiny) and 640 (small) once we have eval data for them
+
+    const int64_t orig_area = static_cast<int64_t>(img.get_size().area());
+
+    size_t  mode_i   = 0;
+    int64_t min_diff = std::numeric_limits<int64_t>::max();
+    for (size_t i = 0; i < std::size(native_resolutions); i++) {
+        const int64_t r    = native_resolutions[i];
+        const int64_t diff = std::abs(orig_area - r * r);
+        if (diff < min_diff) {
+            mode_i   = i;
+            min_diff = diff;
+        }
+    }
+    const int image_size = native_resolutions[mode_i];
+
+    // Aspect-preserving fit-and-pad. Pillow bicubic + PAD_NEAREST for
+    // byte-parity with the upstream deepseek-ai/DeepSeek-OCR HF preprocessor.
+    clip_image_u8 padded;
+    img_tool::resize(img, padded, {image_size, image_size}, RESIZE_ALGO_BICUBIC_PILLOW,
+                     PAD_NEAREST, hparams.image_pad_color);
+    mtmd_image_preproc_out output;
+    output.append_overview(hparams, padded, true);
+    output.grid_x = 0;
+    output.grid_y = 0;
+    // TODO @ngxson : support slicing for DeepSeek-OCR, to do in another PR
+    return output;
+}
+
+//
+// mtmd_image_preprocessor_deepseekocr2
+//
+
+// candidate tile grids (cols, rows) with min_tiles <= cols*rows <= max_tiles
+// sorted by tile count
+std::vector<clip_image_size> mtmd_image_preprocessor_deepseekocr2::get_target_ratios() {
     std::vector<clip_image_size> ratios;
     for (int n = min_tiles; n <= max_tiles; n++) {
         for (int w = 1; w <= n; w++) {
@@ -1134,11 +1171,13 @@ std::vector<clip_image_size> mtmd_image_preprocessor_deepseekocr::get_target_rat
     return ratios;
 }
 
-clip_image_size mtmd_image_preprocessor_deepseekocr::find_closest_aspect_ratio(
+// pick the grid whose aspect ratio is closest to the image
+// on a tie, prefer the larger grid when the image fits
+clip_image_size mtmd_image_preprocessor_deepseekocr2::find_closest_aspect_ratio(
     float                                aspect_ratio,
     const std::vector<clip_image_size> & target_ratios,
     int                                  width,
-    int                                  height) const {
+    int                                  height) {
     float           best_ratio_diff = std::numeric_limits<float>::max();
     clip_image_size best_ratio      = { 1, 1 };
     const float     area            = static_cast<float>(width * height);
@@ -1159,69 +1198,37 @@ clip_image_size mtmd_image_preprocessor_deepseekocr::find_closest_aspect_ratio(
     return best_ratio;
 }
 
-mtmd_image_preproc_out mtmd_image_preprocessor_deepseekocr::preprocess(const clip_image_u8 & img) {
+mtmd_image_preproc_out mtmd_image_preprocessor_deepseekocr2::preprocess(const clip_image_u8 & img) {
+    // emit 768x768 local tiles when the image is larger than a tile in either
+    // dimension, then always a 1024x1024 global view. order: [tiles..., global].
+
     mtmd_image_preproc_out output;
-    int grid_w = 0;
-    int grid_h = 0;
     const auto img_size = img.get_size();
-
-    // global view: aspect-preserving fit-and-pad to base_size
-    clip_image_u8 padded;
-    img_tool::resize(img, padded,
-                    { base_size, base_size },
-                    RESIZE_ALGO_BICUBIC_PILLOW,
-                    PAD_NEAREST,
-                    hparams.image_pad_color);
-    output.append_overview(hparams, padded, true);
-    output.overview.add_viewsep = true;
-
-    // if this condition doesn't hold, the output is overview only, no tiles
     if (img_size.width > tile_size || img_size.height > tile_size) {
         const float           aspect_ratio  = static_cast<float>(img_size.width) / img_size.height;
         const auto            target_ratios = get_target_ratios();
-        const clip_image_size grid =
-            find_closest_aspect_ratio(aspect_ratio, target_ratios, img_size.width, img_size.height);
-        grid_w = grid.width;
-        grid_h = grid.height;
+        const clip_image_size grid          = find_closest_aspect_ratio(aspect_ratio, target_ratios, img_size.width, img_size.height);
 
+        // stretch onto the grid (no aspect preserve), then crop tiles row-major.
         clip_image_u8 refined;
-        img_tool::resize(img, refined, { tile_size * grid_w, tile_size * grid_h }, RESIZE_ALGO_BICUBIC_PILLOW,
-                         PAD_NONE);
+        img_tool::resize(img, refined, { tile_size * grid.width, tile_size * grid.height },
+                         RESIZE_ALGO_BICUBIC_PILLOW, PAD_NONE);
 
-        for (int row = 0; row < grid_h; row++) {
-            if (fuse_row) {
-                // concat all tiles in this row into a single image, along the H axis
-                // output image size: w = tile_size, h = tile_size * grid_w
-                // this is to ensure the whole row is always processed together
-                clip_image_u8 row_img;
-                row_img.set_size({tile_size, tile_size * grid_w}, false);
-                for (int col = 0; col < grid_w; col++) {
-                    for (int py = 0; py < tile_size; py++) {
-                        for (int px = 0; px < tile_size; px++) {
-                            row_img.set_pixel(px, col * tile_size + py,
-                                              refined.get_pixel(col * tile_size + px, row * tile_size + py));
-                        }
-                    }
-                }
-                output.append(hparams, row_img, true);
-            } else {
-                for (int col = 0; col < grid_w; col++) {
-                    clip_image_u8 tile;
-                    img_tool::crop(refined, tile, col * tile_size, row * tile_size, tile_size, tile_size);
-                    output.append(hparams, tile, true);
-                }
+        for (int row = 0; row < grid.height; row++) {
+            for (int col = 0; col < grid.width; col++) {
+                clip_image_u8 tile;
+                img_tool::crop(refined, tile, col * tile_size, row * tile_size, tile_size, tile_size);
+                output.append(hparams, tile, true);
             }
         }
-        if (fuse_row) {
-            grid_w = 1; // each fused row is one image; a single output column
-        }
     }
 
-    LOG_DBG("%s: grid size: %d x %d (%d tiles) + global view\n", __func__, grid_w, grid_h, grid_w * grid_h);
-    LOG_DBG("%s: overview size: %d x %d\n", __func__, padded.get_size().width, padded.get_size().height);
-
-    output.grid_x = grid_w;
-    output.grid_y = grid_h;
+    // global view: aspect-preserving fit-and-pad to base_size.
+    clip_image_u8 padded;
+    img_tool::resize(img, padded, { base_size, base_size }, RESIZE_ALGO_BICUBIC_PILLOW,
+                     PAD_NEAREST, hparams.image_pad_color);
+    output.append_overview(hparams, padded, true);
+    output.overview.add_viewsep = true;
     return output;
 }
 

tools/mtmd/mtmd-image.h

@@ -160,29 +160,29 @@ struct mtmd_image_preprocessor_internvl : mtmd_image_preprocessor_llava_uhd {
     mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override;
 };
 
-// DeepSeek-OCR (v1/v2) global view + optional local tile grid
 struct mtmd_image_preprocessor_deepseekocr : mtmd_image_preprocessor {
-    mtmd_image_preprocessor_deepseekocr(const clip_ctx * ctx)
-        : mtmd_image_preprocessor(ctx),
-            fuse_row(clip_get_projector_type(ctx) == PROJECTOR_TYPE_DEEPSEEKOCR),
-          base_size(hparams.image_size),
-          tile_size(hparams.preproc_tile_size),
-          min_tiles(hparams.preproc_min_tiles),
-          max_tiles(hparams.preproc_max_tiles) {}
+    mtmd_image_preprocessor_deepseekocr(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {}
+    mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override;
+};
+
+// DeepSeek-OCR-2: a 1024x1024 global view, plus InternVL-style 768x768 local
+// tiles when the image is larger than a tile in either dimension.
+struct mtmd_image_preprocessor_deepseekocr2 : mtmd_image_preprocessor {
+    static constexpr int base_size = 1024; // global view
+    static constexpr int tile_size = 768;  // local tile
+    static constexpr int min_tiles = 2;
+    static constexpr int max_tiles = 6;
+
+    mtmd_image_preprocessor_deepseekocr2(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {}
     mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override;
 
 private:
-    bool fuse_row; // v1 fuses a tile-row into one image; v2 keeps tiles separate
-    int base_size; // global view
-    int tile_size; // each tile
-    int min_tiles;
-    int max_tiles;
-
-    std::vector<clip_image_size> get_target_ratios() const;
-    clip_image_size find_closest_aspect_ratio(
-            float aspect_ratio,
-            const std::vector<clip_image_size> & target_ratios,
-            int width, int height) const;
+    static std::vector<clip_image_size> get_target_ratios();
+    static clip_image_size              find_closest_aspect_ratio(
+        float                                aspect_ratio,
+        const std::vector<clip_image_size> & target_ratios,
+        int                                  width,
+        int                                  height);
 };
 
 // custom image preprocessing for Step3VL

tools/mtmd/mtmd.cpp

@@ -618,12 +618,17 @@ struct mtmd_context {
                     image_preproc = std::make_unique<mtmd_image_preprocessor_dyn_size>(ctx_v);
                 } break;
             case PROJECTOR_TYPE_DEEPSEEKOCR:
-            case PROJECTOR_TYPE_DEEPSEEKOCR2:
                 {
                     img_end = "\n"; // prevent empty batch on llama-server
                     image_preproc = std::make_unique<mtmd_image_preprocessor_deepseekocr>(ctx_v);
                     ov_img_first = false;
                 } break;
+            case PROJECTOR_TYPE_DEEPSEEKOCR2:
+                {
+                    img_end = "\n"; // prevent empty batch on llama-server
+                    image_preproc = std::make_unique<mtmd_image_preprocessor_deepseekocr2>(ctx_v);
+                    ov_img_first = false;
+                } break;
             case PROJECTOR_TYPE_HUNYUANVL:
                 {
                     // note: these use fullwidth ｜ (U+FF5C) and ▁ (U+2581) to match the tokenizer vocabulary
@@ -1127,7 +1132,6 @@ struct mtmd_tokenizer {
 
                 // add slices (or tiles)
                 if (!chunks.empty()) {
-                    LOG_DBG("%s: adding %d slices (%d rows x %d cols)\n", __func__, (int)chunks.size(), n_row, n_col);
                     GGML_ASSERT((int)chunks.size() == n_row * n_col);
                     add_text(ctx->tok_slices_start);
                     for (int y = 0; y < n_row; y++) {
@@ -1170,6 +1174,7 @@ struct mtmd_tokenizer {
                     cur.entries.emplace_back(std::move(ov_chunk));
                     add_text(ctx->tok_ov_img_end);
                 }
+
             } else {
 
                 if (preproc_out.entries.size() == 0) {

tools/mtmd/tests/test-deepseek-ocr.py

@@ -29,15 +29,12 @@ class ModelSpec:
     mmproj_arg: str
     model_default: str
     mmproj_default: str
-    prompt: str = "Free OCR."
+    prompt: str = "Free OCR. "
     n_predict: int = 512
     n_ctx: int | None = None
     # Unlimited-OCR's "document parsing" prompt emits <|det|> grounding markup that
     # the HF reference strips in result.md; drop it before scoring to match.
     strip_grounding: bool = False
-    # v2/Unlimited loop on hard tiles; DRY caps it the way HF's
-    # no_repeat_ngram_size does. v1 scores fine without it.
-    dry: bool = False
 
 
 @dataclass
@@ -72,9 +69,6 @@ MODELS = {
         model_arg="--llama-model-2", mmproj_arg="--mmproj-2",
         model_default="gguf_models/deepseek-ai/deepseek-ocr-2-bf16.gguf",
         mmproj_default="gguf_models/deepseek-ai/mmproj-deepseek-ocr-2-bf16.gguf",
-        # v2 keeps generating past 512 on multi-tile; give it room to match the HF ref.
-        n_predict=2048,
-        dry=True,
     ),
     "unlimited": ModelSpec(
         key="unlimited", label="Unlimited-OCR",
@@ -89,7 +83,6 @@ MODELS = {
         n_predict=4096,
         n_ctx=16384,
         strip_grounding=True,
-        dry=True,
     ),
 }
 
@@ -98,9 +91,7 @@ CASES = [
         model_key="v1", label="single-view scan",
         image="tools/mtmd/test-1.jpeg",
         ground_truth="tools/mtmd/tests/test-1-ground-truth.txt",
-        # Fragile image: the HF ref itself swings ~0.286-0.314 across precision
-        # configs -- hence the wide tol. llama.cpp bf16 ~0.322/63.8.
-        hf_cer=0.3140, hf_chrf=67.57, cer_tol=0.04, chrf_tol=5.0,
+        hf_cer=0.3030, hf_chrf=67.52, cer_tol=0.02, chrf_tol=2.0,
     ),
     TestCase(
         model_key="v2", label="single-view scan",
@@ -112,24 +103,6 @@ CASES = [
         # is one pixel off and lands at ~0.69 instead.
         hf_cer=0.7761, hf_chrf=28.70, cer_tol=0.12, chrf_tol=8.0,
     ),
-    TestCase(
-        model_key="v1", label="multi-tile (dynamic resolution)",
-        image="tools/mtmd/tests/test-1-positive.png",
-        ground_truth="tools/mtmd/tests/test-1-ground-truth.txt",
-        # 429x806 -- 806 > 640 triggers the v1 "Gundam" path: (1,2) grid ->
-        # 2 local 640 tiles + 1 global 1024 view. Regression guard for the
-        # tiling preprocessor -- a broken tile path craters the score.
-        # hf_cer/hf_chrf are HF v1's measured scores -- it reads this clean crop exactly.
-        hf_cer=0.0000, hf_chrf=100.00, cer_tol=0.03, chrf_tol=3.0,
-    ),
-    TestCase(
-        model_key="v2", label="multi-tile (dynamic resolution)",
-        image="tools/mtmd/tests/test-1-positive.png",
-        ground_truth="tools/mtmd/tests/test-1-ground-truth.txt",
-        # 429x806 -- 806 > 768 triggers the v2 path: (1,2) grid ->
-        # 2 local 768 tiles + 1 global 1024 view = 545 image tokens.
-        hf_cer=0.0236, hf_chrf=97.05, cer_tol=0.03, chrf_tol=3.0,
-    ),
     TestCase(
         model_key="unlimited", label="single-view scan",
         image="tools/mtmd/test-1.jpeg",
@@ -207,17 +180,14 @@ def run_mtmd_cli(spec: "ModelSpec", model_path, mmproj_path, image_path, bin_pat
         "--flash-attn", "off",  # match the HF "eager" attention reference
         "--no-warmup",
         "-n", str(spec.n_predict),  # cap loops on hard images (KV would otherwise fill)
-    ]
-    if spec.dry:
         # HF decodes with no_repeat_ngram_size; llama.cpp's analog is DRY.
         # Default DRY breakers include "\n", so they are cleared below.
-        cmd += [
-            "--dry-multiplier", "0.8",
-            "--dry-base", "1.75",
-            "--dry-allowed-length", "2",
-            "--dry-penalty-last-n", "-1",
-            "--dry-sequence-breaker", "none",
-        ]
+        "--dry-multiplier", "0.8",
+        "--dry-base", "1.75",
+        "--dry-allowed-length", "2",
+        "--dry-penalty-last-n", "-1",
+        "--dry-sequence-breaker", "none",
+    ]
     if spec.n_ctx is not None:
         cmd += ["-c", str(spec.n_ctx)]
     logger.debug(f"  command: {' '.join(cmd)}")

tools/server/README-dev.md

@@ -126,15 +126,15 @@ It is opt in via the `X-Conversation-Id` header on `POST /v1/chat/completions`.
 
 The feature lives entirely in `server-stream.{h,cpp}` and rests on three types:
 
-- `stream_session`: a bounded ring buffer (4 MiB cap, oldest bytes drop first) plus a condvar. `append` pushes raw SSE bytes, `read_from` drains from any offset and blocks for live bytes or finalize, `finalize` wakes readers, `cancel` sets the flag the producer polls. One conv maps to at most one live session.
+- `stream_session`: a bounded ring buffer (4 MiB cap, oldest bytes drop first) plus a condvar. `append` pushes raw SSE bytes, `read_from` drains from any offset and blocks for live bytes or finalize, `finalize` wakes readers, `cancel` stops the producer. One conv maps to at most one live session.
 - `stream_session_manager`: a file-static singleton (`g_stream_sessions`) inside `server-stream.cpp`, owns all sessions keyed by conv id, enforces the one conv one session invariant via `create_or_replace`, and runs a GC thread that drops completed sessions past their TTL. Exposed to main only through `server_stream_session_manager_start/stop`.
 - `stream_pipe_producer` / `stream_pipe_consumer`: the write and read ends. The producer owns the session lifetime and finalizes it on destruction; the consumer is read only and never finalizes, so a reader detaching cannot kill a running generation.
 
-The implementation is hidden in `server-stream.cpp` (pimpl). The header exposes only the route handler factories, the `server_res_spipe` response base, `server_stream_conv_id_from_headers` and the GC lifecycle; the session, manager, consumer and the `server_stream_create_spipe` factory stay in the `.cpp`.
+The implementation is hidden in `server-stream.cpp` (pimpl). The header exposes only the route handler factories, `server_stream_session_attach_pipe`, `server_stream_aware_should_stop`, `server_stream_conv_id_from_headers` and the GC lifecycle; the session, manager and consumer types stay in the `.cpp`.
 
-Producer side: `server_res_generator` extends `server_res_spipe`, which keeps all spipe logic out of the generic `server_http_res`. `set_req` attaches a producer when the header is present, and the wrapped `next` tees each chunk into the ring before the socket, so a chunk lost to a dead wire is already buffered. While attached, `should_stop` ignores peer disconnect: only a `DELETE` stops generation. On an early peer drop, `on_complete` drains the tail into the ring on the http worker.
+Producer side: `server_res_generator` attaches a producer pipe when the header is present. The HTTP content provider mirrors every chunk into the ring before writing it to the socket. While a pipe is attached, `server_stream_aware_should_stop` ignores peer disconnect, so a dropped socket does not stop generation: only an explicit `DELETE` does. When the peer leaves early, `on_complete` calls `close()`, which drains the rest of the generation into the ring on the http worker.
 
-Lifetime safety: the session holds no back reference to the response, so `spipe` is a plain `unique_ptr` touched only by the http worker. `cancel` raises an atomic the producer polls; the producer finalizes the session from its destructor, which also runs `~server_response_reader::stop()` to cancel the generation at the queue level. A `DELETE` stops work by raising the flag and letting the worker unwind.
+Lifetime safety: the producer pipe holds a shared `alive` flag also captured by the session cancel hook. `~server_res_generator` calls `cleanup()` to clear that hook while the reader is still alive, so a `cancel` arriving during teardown can never call `stop()` on a freed response. This ordering is the most fragile part of the feature: finalizing or destroying the producer before `cleanup()` runs reintroduces a use after free.
 
 Consumer side: `GET /v1/stream/<conv_id>?from=N` opens a `text/event-stream` that replays buffered bytes from offset `N` and blocks for live bytes, so the browser reattaches like a fresh EventSource. An offset below the dropped prefix returns 400.
 
@@ -235,29 +235,6 @@ That requires `JSON.stringify` when formatted to message content:
 }
 ```
 
-Set `stream: true` in the request body to stream a tool's output as it runs, instead of waiting for it to finish. Only certain tools accept this (for ex. `exec_shell_command`);
-returns 404 if tool doesn't support it.
-
-Response is SSE stream, one `data: <json>` line per chunk:
-
-```json
-{"chunk": "hello\n"}
-```
-
-followed by a final event once the tool returns:
-
-```json
-{"done": true}
-```
-
-or, if `invoke()` threw:
-
-```json
-{"done": true, "error": "..."}
-```
-
-There is no `[DONE]` sentinel (unlike `/chat/completions`), the stream ends after the `done`
-
 ### Router mode: how child <--> router communicates
 
 Upon spawning a new child process using `subprocess`, both child and router listen to the stdout/stderr (combined)

tools/server/server-context.cpp

@@ -3979,9 +3979,11 @@ server_context_meta server_context::get_meta() const {
     };
 }
 
+
+
 // generator-like API for HTTP response generation
 // may have bypass_sleep = true if the task does not use ctx_server
-struct server_res_generator : server_res_spipe {
+struct server_res_generator : server_http_res {
     server_response_reader rd;
     server_res_generator(server_queue & queue_tasks, server_response & queue_results, int sleep_idle_seconds, bool bypass_sleep = false)
             : rd(queue_tasks, queue_results, HTTP_POLLING_SECONDS) {
@@ -3991,6 +3993,15 @@ struct server_res_generator : server_res_spipe {
             queue_tasks.wait_until_no_sleep();
         }
     }
+    ~server_res_generator() override {
+        // cleanup() must run while rd is still alive (rd is destroyed after this body returns)
+        if (spipe) {
+            spipe->cleanup();
+        }
+    }
+    void stop() override {
+        rd.stop();
+    }
     void ok(const json & response_data) {
         status = 200;
         data = safe_json_to_str(response_data);
@@ -4028,8 +4039,6 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
     auto & rd = res->rd;
     auto & params = this->params;
 
-    res->set_req(&req); // will also set spipe if needed
-
     int32_t sse_ping_interval = params.sse_ping_interval;
 
     try {
@@ -4172,7 +4181,7 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
         }
         res->status = 200;
         res->content_type = "text/event-stream";
-        res->set_next([res_this = res.get(), res_type, sse_ping_interval](std::string & output) -> bool {
+        res->next = [res_this = res.get(), res_type, sse_ping_interval, &req](std::string & output) -> bool {
             static auto format_error = [](task_response_type res_type, const json & res_json) {
                 if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) {
                     return format_anthropic_sse({
@@ -4184,9 +4193,7 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
                 }
             };
 
-            auto effective_should_stop = [&res_this]() {
-                return res_this->should_stop();
-            };
+            auto effective_should_stop = server_stream_aware_should_stop(res_this, req.should_stop);
 
             try {
                 if (effective_should_stop()) {
@@ -4277,9 +4284,13 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
                 // terminate on exception
                 return false;
             }
-        });
+        };
     }
 
+    // attach a producer pipe to the response when X-Conversation-Id is present.
+    // the pipe mirrors SSE chunks into the ring buffer and wires up the cancel hook.
+    server_stream_session_attach_pipe(*res, req.headers);
+
     return res;
 }
 

tools/server/server-http.cpp

@@ -1,6 +1,7 @@
 #include "common.h"
 #include "http.h"
 #include "server-http.h"
+#include "server-stream.h"
 #include "server-common.h"
 #include "ui.h"
 
@@ -174,15 +175,6 @@ bool server_http_context::init(const common_params & params) {
     // Middlewares
     //
 
-    // Frontend paths - all embedded UI assets
-    static const std::unordered_set<std::string> frontend_paths = []() {
-        std::unordered_set<std::string> paths { "/" };
-        for (const llama_ui_asset & a : llama_ui_get_assets()) {
-            paths.insert("/" + a.name);
-        }
-        return paths;
-    }();
-
     // Public endpoints - API routes plus all embedded UI assets
     static const std::unordered_set<std::string> get_public_endpoints = []() {
         std::unordered_set<std::string> endpoints {
@@ -190,8 +182,11 @@ bool server_http_context::init(const common_params & params) {
             "/v1/health",
             "/models",
             "/v1/models",
+            "/",
         };
-        endpoints.insert(frontend_paths.begin(), frontend_paths.end());
+        for (const llama_ui_asset & a : llama_ui_get_assets()) {
+            endpoints.insert("/" + a.name);
+        }
         return endpoints;
     }();
 
@@ -244,9 +239,18 @@ bool server_http_context::init(const common_params & params) {
 
     auto middleware_server_state = [this](const httplib::Request & req, httplib::Response & res) {
         if (!is_ready.load()) {
-            if (frontend_paths.count(req.path)) {
-                return true; // frontend asset, allow it to load and show "loading"
+#if defined(LLAMA_UI_HAS_ASSETS)
+            if (const auto tmp = string_split<std::string>(req.path, '.');
+                req.path == "/" || (!tmp.empty() && tmp.back() == "html")) {
+                if (const llama_ui_asset * a = llama_ui_find_asset("loading.html")) {
+                    res.status = 503;
+                    res.set_content(reinterpret_cast<const char*>(a->data), a->size, "text/html; charset=utf-8");
+                    return false;
+                }
             }
+#else
+            (void)req;
+#endif
             // no endpoints are allowed to be accessed when the server is not ready
             // this is to prevent any data races or inconsistent states
             res.status = 503;
@@ -529,20 +533,33 @@ static void process_handler_response(server_http_req_ptr && request, server_http
             std::string chunk;
             const bool has_next = response->next(chunk);
             if (!chunk.empty()) {
+                // mirror into the ring buffer first, the session must reflect every SSE chunk
+                // whether or not the wire write below succeeds
+                if (response->spipe) {
+                    response->spipe->write(chunk.data(), chunk.size());
+                }
                 if (!sink.write(chunk.data(), chunk.size())) {
+                    // peer is gone, stop the wire path here
                     return false;
                 }
                 SRV_DBG("http: streamed chunk: %s\n", chunk.c_str());
             }
             if (!has_next) {
+                // producer reached its natural end on the wire, a later close() skips the drain
+                if (response->spipe) {
+                    response->spipe->done();
+                }
                 sink.done();
                 SRV_DBG("%s", "http: stream ended\n");
             }
             return has_next;
         };
         const auto on_complete = [request = q_ptr, response = r_ptr](bool) mutable {
-            response->on_complete();
-            response.reset();
+            // on a dropped peer, close() drains the rest of the generation into the ring buffer
+            if (response->spipe) {
+                response->spipe->close();
+            }
+            response.reset(); // spipe destructor finalizes the session if attached
             request.reset();
         };
         res.set_chunked_content_provider(content_type, chunked_content_provider, on_complete);
@@ -550,7 +567,6 @@ static void process_handler_response(server_http_req_ptr && request, server_http
         res.status = response->status;
         set_headers(res, response->headers);
         res.set_content(response->data, response->content_type);
-        response->on_complete();
     }
 }
 

tools/server/server-http.h

@@ -11,6 +11,7 @@
 #include <unordered_map>
 
 struct common_params;
+struct stream_pipe_producer; // defined in server-stream.h
 
 // generator-like API for HTTP response generation
 // this object response with one of the 2 modes:
@@ -24,13 +25,19 @@ struct server_http_res {
     std::string data;
     std::map<std::string, std::string> headers;
 
+    // if set, the stream survives a client disconnect: the producer pipe keeps draining into the
+    // ring buffer and finalizes the session on destruction, so no explicit on_stream_end is needed.
+    // shared_ptr (not unique_ptr) so the forward-declared type is safe to delete here.
+    std::shared_ptr<stream_pipe_producer> spipe;
+
     std::function<bool(std::string &)> next = nullptr;
     bool is_stream() const {
         return next != nullptr;
     }
 
-    // fired before req and res are destroyed
-    virtual void on_complete() {}
+    // called when the session is cancelled (e.g. DELETE /v1/stream/<conv_id>).
+    // server_res_generator overrides this to stop its reader; the default is a no-op.
+    virtual void stop() {}
 
     virtual ~server_http_res() = default;
 };

tools/server/server-schema.cpp

@@ -568,16 +568,10 @@ static void handle_with_catch(const char * name, std::function<void()> func) {
     }
 }
 
-// treat a null value as absent so clients can send null to request the server default
-static bool has_value(const json & data, const char * n) {
-    auto it = data.find(n);
-    return it != data.end() && !it->is_null();
-}
-
 template <typename T>
 void field_num<T>::eval(field_eval_context & ctx, const json & data) {
     for (const auto & n : name) {
-        if (has_value(data, n)) {
+        if (data.contains(n)) {
             handle_with_catch(n, [&]() {
                 if (custom_handler) {
                 custom_handler(ctx, data);
@@ -599,7 +593,7 @@ void field_num<T>::eval(field_eval_context & ctx, const json & data) {
 void field_str::eval(field_eval_context & ctx, const json & data) {
     GGML_ASSERT(custom_handler);
     for (const auto & n : name) {
-        if (has_value(data, n)) {
+        if (data.contains(n)) {
             handle_with_catch(n, [&]() {
                 custom_handler(ctx, data);
             });
@@ -610,7 +604,7 @@ void field_str::eval(field_eval_context & ctx, const json & data) {
 
 void field_bool::eval(field_eval_context & ctx, const json & data) {
     for (const auto & n : name) {
-        if (has_value(data, n)) {
+        if (data.contains(n)) {
             handle_with_catch(n, [&]() {
                 if (custom_handler) {
                     custom_handler(ctx, data);
@@ -626,7 +620,7 @@ void field_bool::eval(field_eval_context & ctx, const json & data) {
 void field_json::eval(field_eval_context & ctx, const json & data) {
     GGML_ASSERT(custom_handler);
     for (const auto & n : name) {
-        if (has_value(data, n)) {
+        if (data.contains(n)) {
             handle_with_catch(n, [&]() {
                 custom_handler(ctx, data);
             });

tools/server/server-stream.cpp

@@ -96,6 +96,8 @@ struct stream_session {
     size_t  dropped_prefix() const; // bytes evicted from the front due to cap
     int64_t completed_at() const;   // 0 while alive, unix seconds after finalize
 
+    void set_stop_producer(std::function<void()> fn);
+
     void cancel();
 
 private:
@@ -107,6 +109,7 @@ private:
     bool                    done;
     std::atomic<bool>       cancelled; // polled lock-free by the should_stop closure, no mu
     int64_t                 completed_ts;
+    std::function<void()>   stop_producer;
 };
 stream_session::stream_session(std::string conversation_id_, size_t max_bytes_)
     : conversation_id(std::move(conversation_id_))
@@ -214,10 +217,26 @@ int64_t stream_session::completed_at() const {
     return completed_ts;
 }
 
+void stream_session::set_stop_producer(std::function<void()> fn) {
+    std::lock_guard<std::mutex> lock(mu);
+    stop_producer = std::move(fn);
+}
+
 void stream_session::cancel() {
-    // the should_stop closure on both the producer and any HTTP reader polls is_cancelled()
-    // so flipping this is the only signal needed to unwind both sides
+    // flip cancelled first so the producer-side server_stream_aware_should_stop can break out of the
+    // recv() wait even if remove_waiting_task_ids does not notify the condvar (the cancel task
+    // posted by rd.stop() will eventually notify, but we do not want to depend on that timing)
     cancelled.store(true, std::memory_order_release);
+    // copy the hook under the lock then invoke outside, the producer side may grab queue locks
+    // and we do not want to hold our mu across that path
+    std::function<void()> fn;
+    {
+        std::lock_guard<std::mutex> lock(mu);
+        fn = stop_producer;
+    }
+    if (fn) {
+        fn();
+    }
 }
 
 bool stream_session::is_cancelled() const {
@@ -306,10 +325,8 @@ void stream_session_manager::evict_and_cancel(const std::string & conversation_i
         s = it->second;
         sessions.erase(it);
     }
-    // cancel first so the producer's on_complete() drain loop and any pending HTTP reader
-    // observe is_cancelled() and stop pulling further output, then finalize to wake readers
-    // blocked in read_from(). note: this does not interrupt the underlying generation itself,
-    // which keeps running to its own natural stop condition (EOS/max_tokens)
+    // signal the producer side first so the inference is cancelled at the queue level,
+    // then finalize, which wakes any pending HTTP reader and lets the drain exit naturally
     s->cancel();
     s->finalize();
 }
@@ -414,15 +431,65 @@ stream_pipe_producer::stream_pipe_producer(stream_session_ptr session)
 }
 
 stream_pipe_producer::~stream_pipe_producer() {
+    cleanup();
     session_->finalize();
 }
 
+void stream_pipe_producer::cleanup() {
+    if (!alive_) {
+        return;
+    }
+    alive_->store(false, std::memory_order_release);
+    session_->set_stop_producer(nullptr);
+    alive_.reset();
+}
+
 bool stream_pipe_producer::write(const char * data, size_t len) {
     return session_->append(data, len);
 }
 
-stream_pipe_producer * stream_pipe_producer::create(stream_session_ptr session) {
-    return new stream_pipe_producer(std::move(session));
+void stream_pipe_producer::done() {
+    done_ = true;
+}
+
+void stream_pipe_producer::close() {
+    // httplib bails its content provider the moment is_peer_alive() goes false, so pump the rest
+    // of the generation into the ring buffer here. a DELETE flips is_cancelled and cuts it short
+    if (done_ || session_->is_cancelled()) {
+        SRV_TRC("stream_pipe close: skip drain (done=%d cancelled=%d) conv=%s\n",
+                done_ ? 1 : 0, session_->is_cancelled() ? 1 : 0, session_->conversation_id.c_str());
+        return;
+    }
+    SRV_TRC("stream_pipe close: draining conv=%s\n", session_->conversation_id.c_str());
+    size_t drained = 0;
+    std::string chunk;
+    while (true) {
+        chunk.clear();
+        bool has_next = res_->next(chunk);
+        if (!chunk.empty()) {
+            write(chunk.data(), chunk.size());
+            drained += chunk.size();
+        }
+        if (!has_next) {
+            break;
+        }
+    }
+    SRV_TRC("stream_pipe close: drain ended conv=%s bytes=%zu\n", session_->conversation_id.c_str(), drained);
+}
+
+std::shared_ptr<stream_pipe_producer> stream_pipe_producer::create(stream_session_ptr session,
+                                                                   server_http_res & res) {
+    auto alive = std::make_shared<std::atomic<bool>>(true);
+    auto * res_ptr = &res;
+    session->set_stop_producer([alive, res_ptr]() {
+        if (alive->load(std::memory_order_acquire)) {
+            res_ptr->stop();
+        }
+    });
+    auto pipe = std::shared_ptr<stream_pipe_producer>(new stream_pipe_producer(std::move(session)));
+    pipe->alive_ = std::move(alive);
+    pipe->res_   = res_ptr;
+    return pipe;
 }
 
 // stream_pipe_consumer
@@ -594,68 +661,21 @@ std::string server_stream_conv_id_from_headers(const std::map<std::string, std::
     return std::string();
 }
 
-static stream_pipe_producer * server_stream_create_spipe(const std::map<std::string, std::string> & headers) {
+void server_stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers) {
     std::string conversation_id = server_stream_conv_id_from_headers(headers);
     SRV_TRC("conv_id=%s (empty=%d)\n", conversation_id.c_str(), conversation_id.empty() ? 1 : 0);
     if (conversation_id.empty()) {
-        return nullptr;
-    }
-    auto session = g_stream_sessions.create_or_replace(conversation_id);
-    return stream_pipe_producer::create(session);
-}
-
-//
-// server_res_spipe
-//
-
-void server_res_spipe::set_req(const server_http_req * req) {
-    this->req = req;
-    // optionally attach spipe to the response when X-Conversation-Id is present
-    spipe.reset(server_stream_create_spipe(req->headers));
-}
-
-bool server_res_spipe::conn_alive() {
-    GGML_ASSERT(req != nullptr);
-    return !req->should_stop();
-}
-
-bool server_res_spipe::should_stop() {
-    if (spipe) {
-        // note: if DELETE /v1/stream/<conv_id> is called, is_cancelled() will be true
-        return spipe->is_cancelled();
-    } else {
-        return !conn_alive();
-    }
-}
-
-void server_res_spipe::on_complete() {
-    if (!spipe || next_finished) {
         return;
     }
-    std::string chunk;
-    while (!spipe->is_cancelled()) {
-        chunk.clear();
-        bool has_next = next_orig(chunk);
-        if (!chunk.empty()) {
-            spipe->write(chunk.data(), chunk.size());
-        }
-        if (!has_next) {
-            break;
-        }
-    }
+    auto session = g_stream_sessions.create_or_replace(conversation_id);
+    res.spipe = stream_pipe_producer::create(session, res);
 }
 
-void server_res_spipe::set_next(std::function<bool(std::string &)> next_fn) {
-    next_orig = std::move(next_fn);
-    next = [this](std::string & out) {
-        bool has_next = next_orig(out);
-        if (spipe) {
-            // if spipe is set, tee-style pipe input to both HTTP and spipe
-            spipe->write(out.data(), out.size());
+std::function<bool()> server_stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback) {
+    return [res, fallback = std::move(fallback)]() -> bool {
+        if (res->spipe) {
+            return res->spipe->is_cancelled();
         }
-        if (!has_next) {
-            next_finished = true;
-        }
-        return has_next;
+        return fallback();
     };
 }

tools/server/server-stream.h

@@ -30,15 +30,36 @@ protected:
 
 // producer end: writes chunks into the ring buffer and owns the session lifetime, finalizing it
 // on destruction.
+//
+// lifetime safety: holds a shared_ptr<atomic<bool>> alive also captured by the session's
+// stop_producer hook. cleanup() sets alive=false and clears the hook; it must run while the
+// response the hook calls stop() on is still alive. ~server_res_generator() does this explicitly.
 struct stream_pipe_producer : stream_pipe {
     ~stream_pipe_producer() override;
 
     bool write(const char * data, size_t len);
 
-    static stream_pipe_producer * create(stream_session_ptr session);
+    // mark the natural end on the wire so a later close() is a no-op
+    void done();
+
+    // on a peer drop, pump the response next() into the ring buffer until done. runs on the http
+    // worker from on_complete, no-op after done() or cancel
+    void close();
+
+    // disarm the stop hook and drop the alive guard, must run while the response the hook
+    // references is still alive. idempotent, the destructor calls it too
+    void cleanup();
+
+    // res.stop() is invoked when the session is cancelled, the alive guard ensures stop() is not
+    // called after cleanup() has run
+    static std::shared_ptr<stream_pipe_producer> create(stream_session_ptr session, server_http_res & res);
 
 private:
     explicit stream_pipe_producer(stream_session_ptr session);
+
+    bool                                done_ = false;
+    std::shared_ptr<std::atomic<bool>>  alive_;
+    server_http_res *                   res_ = nullptr;
 };
 
 void server_stream_session_manager_start();
@@ -52,22 +73,10 @@ server_http_context::handler_t server_stream_make_delete_handler();
 // extract the X-Conversation-Id header value (case-insensitive), empty when absent
 std::string server_stream_conv_id_from_headers(const std::map<std::string, std::string> & headers);
 
-// implement tee-style pipe (spipe) for "stream replay" functionality
-struct server_res_spipe : server_http_res {
-private:
-    // if set, the stream survives a client disconnect:
-    // connection kept alive, output is forwarded to spipe and reuse later
-    std::unique_ptr<stream_pipe_producer> spipe;
-    // if spipe is set, use this next_orig to implement tee-style pipe
-    std::function<bool(std::string &)> next_orig;
-    const server_http_req * req = nullptr;
-    // set once next_orig reports no more data, so on_complete() doesn't re-drain a finished stream
-    bool next_finished = false;
+// on an X-Conversation-Id header, create or replace the session and attach a producer pipe to res
+void server_stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers);
 
-public:
-    void set_req(const server_http_req * req);
-    bool conn_alive();
-    bool should_stop();
-    void on_complete() override;
-    void set_next(std::function<bool(std::string &)> next_fn);
-};
+// should_stop closure that ignores peer disconnect when a pipe is attached, so only an explicit
+// DELETE stops the producer and generation keeps flowing into the ring buffer. without a pipe it
+// delegates to fallback, the legacy non-resumable flow
+std::function<bool()> server_stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback);

tools/server/server-tools.cpp

@@ -12,7 +12,6 @@
 #include <climits>
 #include <algorithm>
 #include <unordered_set>
-#include <functional>
 
 namespace fs = std::filesystem;
 
@@ -52,13 +51,7 @@ public:
     virtual bool write_file(const std::string & path, const std::string & content) const = 0;
     // paths relative to `base`, '/'-separated; sets `err` if `base` isn't a directory
     virtual std::vector<std::string> list_files(const std::string & base, std::string & err) const = 0;
-    // on_chunk, if set, is called with each chunk of output as it is read (before truncation cuts in);
-    // returning false terminates the process early (e.g. the client disconnected)
-    virtual exec_result run(
-            const std::vector<std::string> & args,
-            size_t max_output,
-            int timeout_secs,
-            const std::function<bool(const std::string &)> & on_chunk = nullptr) const = 0;
+    virtual exec_result run(const std::vector<std::string> & args, size_t max_output, int timeout_secs) const = 0;
 };
 
 class tools_io_basic : public tools_io {
@@ -130,11 +123,7 @@ public:
         return list_files_fallback(base);
     }
 
-    exec_result run(
-            const std::vector<std::string> & args,
-            size_t max_output,
-            int timeout_secs,
-            const std::function<bool(const std::string &)> & on_chunk = nullptr) const override {
+    exec_result run(const std::vector<std::string> & args, size_t max_output, int timeout_secs) const override {
         exec_result res;
 
         subprocess_s proc;
@@ -175,14 +164,8 @@ public:
                     size_t len = strlen(buf);
                     if (output.size() + len <= max_output) {
                         output.append(buf, len);
-                        if (on_chunk && !on_chunk(std::string(buf, len))) {
-                            subprocess_terminate(&proc);
-                            break;
-                        }
                     } else {
-                        size_t remaining = max_output - output.size();
-                        output.append(buf, remaining);
-                        if (on_chunk && remaining > 0) on_chunk(std::string(buf, remaining));
+                        output.append(buf, max_output - output.size());
                         truncated = true;
                     }
                 }
@@ -304,7 +287,7 @@ struct server_tool_read_file : server_tool {
         };
     }
 
-    json invoke(json params, server_tool::stream *) const override {
+    json invoke(json params) const override {
         std::string path  = params.at("path").get<std::string>();
         int  start_line   = json_value(params, "start_line", 1);
         int  end_line     = json_value(params, "end_line",  -1); // -1 = no limit
@@ -393,7 +376,7 @@ struct server_tool_file_glob_search : server_tool {
         };
     }
 
-    json invoke(json params, server_tool::stream *) const override {
+    json invoke(json params) const override {
         std::string base    = params.at("path").get<std::string>();
         std::string include = json_value(params, "include", std::string("**"));
         std::string exclude = json_value(params, "exclude", std::string(""));
@@ -474,7 +457,7 @@ struct server_tool_grep_search : server_tool {
         };
     }
 
-    json invoke(json params, server_tool::stream *) const override {
+    json invoke(json params) const override {
         std::string path        = params.at("path").get<std::string>();
         std::string pat_str     = params.at("pattern").get<std::string>();
         std::string include     = json_value(params, "include", std::string("**"));
@@ -594,7 +577,6 @@ struct server_tool_exec_shell_command : server_tool {
         name = "exec_shell_command";
         display_name = "Execute shell command";
         permission_write = true;
-        support_stream = true;
     }
 
     json get_definition() const override {
@@ -616,7 +598,7 @@ struct server_tool_exec_shell_command : server_tool {
         };
     }
 
-    json invoke(json params, server_tool::stream * st) const override {
+    json invoke(json params) const override {
         std::string command   = params.at("command").get<std::string>();
         int    timeout        = json_value(params, "timeout",         10);
         size_t max_output     = (size_t) json_value(params, "max_output_size", (int) SERVER_TOOL_EXEC_SHELL_COMMAND_MAX_OUTPUT_SIZE);
@@ -630,24 +612,7 @@ struct server_tool_exec_shell_command : server_tool {
         std::vector<std::string> args = {"sh", "-c", command};
 #endif
 
-        auto io = make_tools_io(params);
-
-        if (st) {
-            auto res = io->run(args, max_output, timeout, [st](const std::string & chunk) {
-                st->push(chunk);
-                return !st->alive || st->alive();
-            });
-            if (st->alive && !st->alive()) {
-                return json();
-            }
-            std::string tail = string_format("\n[exit code: %d]", res.exit_code);
-            if (res.timed_out) {
-                tail += " [exit due to timed out]";
-            }
-            st->push(tail);
-            return json();
-        }
-
+        auto io  = make_tools_io(params);
         auto res = io->run(args, max_output, timeout);
 
         std::string text_output = res.output;
@@ -689,7 +654,7 @@ struct server_tool_write_file : server_tool {
         };
     }
 
-    json invoke(json params, server_tool::stream *) const override {
+    json invoke(json params) const override {
         std::string path    = params.at("path").get<std::string>();
         std::string content = params.at("content").get<std::string>();
 
@@ -745,7 +710,7 @@ struct server_tool_edit_file : server_tool {
         };
     }
 
-    json invoke(json params, server_tool::stream *) const override {
+    json invoke(json params) const override {
         std::string path = params.at("path").get<std::string>();
         const json & edits_json = params.at("edits");
 
@@ -1053,7 +1018,7 @@ struct server_tool_get_datetime : server_tool {
         };
     }
 
-    json invoke(json, server_tool::stream *) const override {
+    json invoke(json) const override {
         auto now = std::chrono::system_clock::now();
         auto time = std::chrono::system_clock::to_time_t(now);
 
@@ -1061,59 +1026,6 @@ struct server_tool_get_datetime : server_tool {
     }
 };
 
-struct server_tool_stream_result : server_task_result {
-    std::string chunk;
-    bool done = false;
-    std::string error_msg;
-
-    json to_json() override {
-        if (!done) {
-            return {{"chunk", chunk}};
-        } else {
-            json result = {{"done", true}};
-            if (!error_msg.empty()) {
-                result["error"] = error_msg;
-            }
-            return result;
-        }
-    }
-};
-
-void server_tool::stream::push(const std::string & chunk) {
-    if (chunk.empty()) return;
-    auto r = std::make_unique<server_tool_stream_result>();
-    r->id    = id;
-    r->chunk = chunk;
-    qr.send(std::move(r));
-}
-
-struct server_tools_res : server_http_res {
-    std::thread worker;
-    server_response * qr = nullptr; // set only for streaming responses
-    int id = -1;
-
-    ~server_tools_res() override {
-        if (worker.joinable()) {
-            worker.join();
-        }
-        if (qr) {
-            qr->remove_waiting_task_id(id);
-        }
-    }
-};
-
-static server_tool & find_tool(std::vector<std::unique_ptr<server_tool>> & tools, const std::string & name, bool require_stream) {
-    for (auto & t : tools) {
-        if (t->name == name) {
-            if (require_stream && !t->support_stream) {
-                throw std::invalid_argument(string_format("tool \"%s\" does not support stream = true", name.c_str()));
-            }
-            return *t;
-        }
-    }
-    throw std::invalid_argument(string_format("unknown tool \"%s\"", name.c_str()));
-}
-
 //
 // public API
 //
@@ -1178,63 +1090,16 @@ void server_tools::setup(const std::vector<std::string> & enabled_tools) {
     };
 
     handle_post = [this](const server_http_req & req) -> server_http_res_ptr {
-        auto res = std::make_unique<server_tools_res>();
+        auto res = std::make_unique<server_http_res>();
         try {
             json body = json::parse(req.body);
             std::string tool_name = body.at("tool").get<std::string>();
             json params = body.value("params", json::object());
-            bool stream = body.value("stream", false);
-
-            server_tool & tool = find_tool(tools, tool_name, stream);
-
-            if (stream) {
-                int id = res_id.fetch_add(1);
-                queue_res.add_waiting_task_id(id);
-                res->qr = &queue_res;
-                res->id = id;
-
-                res->worker = std::thread([this, id, &req, &tool, params]() mutable {
-                    server_tool::stream st{queue_res, id, [&req]() {
-                        return !req.should_stop();
-                    }};
-
-                    auto done = std::make_unique<server_tool_stream_result>();
-                    try {
-                        tool.invoke(params, &st);
-                    } catch (const std::exception & e) {
-                        done->error_msg = e.what();
-                    } catch (...) {
-                        done->error_msg = "An unknown error occurred";
-                    }
-                    done->id    = st.id;
-                    done->done  = true;
-                    st.qr.send(std::move(done));
-                });
-
-                res->content_type = "text/event-stream";
-                res->status = 200;
-                res->next   = [this, id](std::string & output) -> bool {
-                    auto result = queue_res.recv(id);
-                    auto * r = dynamic_cast<server_tool_stream_result *>(result.get());
-                    GGML_ASSERT(r != nullptr);
-                    output = "data: " + safe_json_to_str(r->to_json()) + "\n\n";
-                    if (r->done) {
-                        queue_res.remove_waiting_task_id(id);
-                        return false;
-                    }
-                    return true;
-                };
-            } else {
-                json result = tool.invoke(params, nullptr);
-                res->status = 200;
-                res->data   = safe_json_to_str(result);
-            }
+            json result = invoke(tool_name, params);
+            res->data   = safe_json_to_str(result);
         } catch (const json::exception & e) {
             res->status = 400;
             res->data   = safe_json_to_str(format_error_response(e.what(), ERROR_TYPE_INVALID_REQUEST));
-        } catch (const std::invalid_argument & e) {
-            res->status = 404;
-            res->data   = safe_json_to_str(format_error_response(e.what(), ERROR_TYPE_INVALID_REQUEST));
         } catch (const std::exception & e) {
             SRV_ERR("got exception: %s\n", e.what());
             res->status = 500;
@@ -1243,3 +1108,12 @@ void server_tools::setup(const std::vector<std::string> & enabled_tools) {
         return res;
     };
 }
+
+json server_tools::invoke(const std::string & name, const json & params) {
+    for (auto & t : tools) {
+        if (t->name == name) {
+            return t->invoke(params);
+        }
+    }
+    return {{"error", "unknown tool: " + name}};
+}

tools/server/server-tools.h

@@ -2,27 +2,15 @@
 
 #include "server-common.h"
 #include "server-http.h"
-#include "server-queue.h"
-
-#include <atomic>
-#include <functional>
 
 struct server_tool {
     std::string name;
     std::string display_name;
     bool permission_write = false;
-    bool support_stream = false; // if true, output can be streamed
 
     virtual ~server_tool() = default;
     virtual json get_definition() const = 0;
-
-    struct stream {
-        server_response & qr;
-        int id;
-        std::function<bool()> alive;
-        void push(const std::string & chunk);
-    };
-    virtual json invoke(json params, stream * st = nullptr) const = 0;
+    virtual json invoke(json params) const = 0;
 
     json to_json() const;
 };
@@ -30,11 +18,8 @@ struct server_tool {
 struct server_tools {
     std::vector<std::unique_ptr<server_tool>> tools;
 
-    // for streaming
-    server_response queue_res;
-    std::atomic<int> res_id{0};
-
     void setup(const std::vector<std::string> & enabled_tools);
+    json invoke(const std::string & name, const json & params);
 
     server_http_context::handler_t handle_get;
     server_http_context::handler_t handle_post;

tools/server/tests/unit/test_tools_builtin.py

@@ -105,24 +105,6 @@ def test_tools_builtin_edit_file_rejects_non_unique_old_text():
             os.remove(log_path)
 
 
-def test_tools_builtin_exec_shell_command_stream():
-    global server
-    server.start()
-
-    events = list(server.make_stream_request("POST", "/tools", data={
-        "tool": "exec_shell_command",
-        "params": {"command": "echo hello"},
-        "stream": True,
-    }))
-
-    assert len(events) >= 2
-    assert events[-1]["done"] is True
-    assert not events[-1].get("error")
-    chunks = "".join(e["chunk"] for e in events[:-1])
-    assert "hello" in chunks
-    assert "[exit code: 0]" in chunks
-
-
 def test_tools_builtin_edit_file_rejects_overlapping_edits():
     global server
     server.start()

tools/ui/embed.cpp

@@ -187,6 +187,7 @@ int main(int argc, char ** argv) {
         struct required_check { const char * label; match_fn match; bool found; };
         required_check checks[] = {
             { "index.html",           exact("index.html"),           false },
+            { "loading.html",         exact("loading.html"),         false },
             { "manifest.webmanifest", exact("manifest.webmanifest"), false },
             { "sw.js",                exact("sw.js"),                false },
             { "build.json",           exact("build.json"),           false },

tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenServerError.svelte

@@ -1,11 +1,10 @@
 <script lang="ts">
-	import { AlertTriangle, Loader2, RefreshCw } from '@lucide/svelte';
+	import { AlertTriangle, RefreshCw } from '@lucide/svelte';
 	import { fadeInView } from '$lib/actions/fade-in-view.svelte';
 	import * as Alert from '$lib/components/ui/alert';
-	import { serverError, serverLoading, serverStatus, serverStore } from '$lib/stores/server.svelte';
+	import { serverError, serverLoading, serverStore } from '$lib/stores/server.svelte';
 
 	let hasError = $derived(!!serverError());
-	let isLoadingModel = $derived(serverStatus() === 503);
 </script>
 
 {#if hasError}
@@ -13,31 +12,23 @@
 		class="pointer-events-auto mx-auto mb-4 max-w-[48rem] px-1"
 		use:fadeInView={{ y: 10, duration: 250 }}
 	>
-		<Alert.Root variant={isLoadingModel ? 'default' : 'destructive'}>
-			{#if isLoadingModel}
-				<Loader2 class="h-4 w-4 animate-spin" />
-			{:else}
-				<AlertTriangle class="h-4 w-4" />
-			{/if}
+		<Alert.Root variant="destructive">
+			<AlertTriangle class="h-4 w-4" />
 
 			<Alert.Title class="flex items-center justify-between">
-				<span>{isLoadingModel ? 'Loading model' : 'Server unavailable'}</span>
+				<span>Server unavailable</span>
 
-				{#if !isLoadingModel}
-					<button
-						onclick={() => serverStore.fetch()}
-						disabled={serverLoading()}
-						class="flex items-center gap-1.5 rounded-lg bg-destructive/20 px-2 py-1 text-xs font-medium hover:bg-destructive/30 disabled:opacity-50"
-					>
-						<RefreshCw class="h-3 w-3 {serverLoading() ? 'animate-spin' : ''}" />
-						{serverLoading() ? 'Retrying...' : 'Retry'}
-					</button>
-				{/if}
+				<button
+					onclick={() => serverStore.fetch()}
+					disabled={serverLoading()}
+					class="flex items-center gap-1.5 rounded-lg bg-destructive/20 px-2 py-1 text-xs font-medium hover:bg-destructive/30 disabled:opacity-50"
+				>
+					<RefreshCw class="h-3 w-3 {serverLoading() ? 'animate-spin' : ''}" />
+					{serverLoading() ? 'Retrying...' : 'Retry'}
+				</button>
 			</Alert.Title>
 
-			{#if !isLoadingModel}
-				<Alert.Description>{serverError()}</Alert.Description>
-			{/if}
+			<Alert.Description>{serverError()}</Alert.Description>
 		</Alert.Root>
 	</div>
 {/if}

tools/ui/src/lib/constants/pwa.ts

@@ -258,6 +258,12 @@ export const GLOB_PATTERNS: string[] = [
 	'**/*.{js,css,html,ico,svg,png,webp,woff,woff2,json,webmanifest}'
 ];
 
+// loading.html is the model loading page served by llama-server itself.
+// The SvelteKit PWA manifest transform strips the html extension from every
+// precache entry to match clean URLs, but loading.html is a plain static asset
+// with no clean URL, so static servers answer 404 and the SW install fails.
+export const GLOB_IGNORES: string[] = ['**/loading.html'];
+
 export const SW_CONFIG = {
 	CHECK_INTERVAL_MS: 60000,
 	UPDATE_FETCH_OPTIONS: {
@@ -311,6 +317,7 @@ export const SVELTEKIT_PWA_OPTIONS: SvelteKitPWAOptions = {
 		// Uses '**/' because SvelteKit outputs files under _app/immutable/
 		// subdirectories.
 		globPatterns: GLOB_PATTERNS,
+		globIgnores: GLOB_IGNORES,
 		maximumFileSizeToCacheInBytes: CACHE_SETTINGS.MAX_FILE_SIZE_BYTES,
 
 		// Prevent @vite-pwa/sveltekit from auto-adding a NavigationRoute by

tools/ui/src/lib/stores/models.svelte.ts

@@ -145,10 +145,6 @@ class ModelsStore {
 	 */
 
 	getModelModalities(modelId: string): ModelModalities | null {
-		if (!isRouterMode() && serverStore.props?.modalities) {
-			return this.buildModalities(serverStore.props.modalities);
-		}
-
 		const model = this.models.find((m) => m.model === modelId || m.id === modelId);
 		if (model?.modalities) {
 			return model.modalities;
@@ -633,12 +629,7 @@ class ModelsStore {
 	}
 
 	findModelByName(modelName: string): ModelOption | null {
-		return (
-			this.models.find(
-				(model) =>
-					model.model === modelName || model.id === modelName || model.aliases?.includes(modelName)
-			) ?? null
-		);
+		return this.models.find((model) => model.model === modelName) ?? null;
 	}
 
 	findModelById(modelId: string): ModelOption | null {

tools/ui/src/lib/stores/server.svelte.ts

@@ -1,8 +1,5 @@
 import { PropsService } from '$lib/services/props.service';
 import { ServerRole } from '$lib/enums';
-import { ApiError } from '$lib/utils/api-fetch';
-
-const LOADING_RETRY_INTERVAL_MS = 1000;
 
 /**
  * serverStore - Server connection state, configuration, and role detection
@@ -32,10 +29,8 @@ class ServerStore {
 	props = $state<ApiLlamaCppServerProps | null>(null);
 	loading = $state(false);
 	error = $state<string | null>(null);
-	status = $state<number | null>(null);
 	role = $state<ServerRole | null>(null);
 	private fetchPromise: Promise<void> | null = null;
-	private retryTimer: ReturnType<typeof setTimeout> | null = null;
 
 	/**
 	 *
@@ -75,43 +70,23 @@ class ServerStore {
 	 *
 	 */
 
-	/**
-	 * @param background - Set by the automatic "still loading" poll. Skips the
-	 * `loading` flag flip so the UI doesn't bounce between the full loading
-	 * splash and the chat screen every retry tick.
-	 */
-	async fetch({ background = false }: { background?: boolean } = {}): Promise<void> {
+	async fetch(): Promise<void> {
 		if (this.fetchPromise) return this.fetchPromise;
 
-		this.clearRetryTimer();
-		if (!background) {
-			this.loading = true;
-		}
-		// Don't clear an existing "still loading" error before a retry -
-		// doing so would unmount/remount the error banner every second.
-		if (this.status !== 503) {
-			this.error = null;
-		}
+		this.loading = true;
+		this.error = null;
 
 		const fetchPromise = (async () => {
 			try {
 				const props = await PropsService.fetch();
 				this.props = props;
 				this.error = null;
-				this.status = null;
 				this.detectRole(props);
 			} catch (error: unknown) {
 				this.error = error instanceof Error ? error.message : String(error);
-				this.status = error instanceof ApiError ? error.status : null;
 				console.error('Error fetching server properties:', error);
-
-				if (this.status === 503) {
-					this.scheduleRetry();
-				}
 			} finally {
-				if (!background) {
-					this.loading = false;
-				}
+				this.loading = false;
 				this.fetchPromise = null;
 			}
 		})();
@@ -121,30 +96,13 @@ class ServerStore {
 	}
 
 	clear(): void {
-		this.clearRetryTimer();
 		this.props = null;
 		this.error = null;
-		this.status = null;
 		this.loading = false;
 		this.role = null;
 		this.fetchPromise = null;
 	}
 
-	private scheduleRetry(): void {
-		if (this.retryTimer) return;
-		this.retryTimer = setTimeout(() => {
-			this.retryTimer = null;
-			this.fetch({ background: true });
-		}, LOADING_RETRY_INTERVAL_MS);
-	}
-
-	private clearRetryTimer(): void {
-		if (this.retryTimer) {
-			clearTimeout(this.retryTimer);
-			this.retryTimer = null;
-		}
-	}
-
 	/**
 	 *
 	 *
@@ -167,7 +125,6 @@ export const serverStore = new ServerStore();
 export const serverProps = () => serverStore.props;
 export const serverLoading = () => serverStore.loading;
 export const serverError = () => serverStore.error;
-export const serverStatus = () => serverStore.status;
 export const serverRole = () => serverStore.role;
 export const defaultParams = () => serverStore.defaultParams;
 export const contextSize = () => serverStore.contextSize;

tools/ui/src/lib/utils/api-fetch.ts

@@ -12,21 +12,6 @@ import { ERROR_MESSAGES, HTTP_CODE_TO_STRING } from '$lib/constants/error';
  * - Base path resolution
  */
 
-/**
- * Error thrown when an API request fails, carrying the HTTP status code
- * so callers can distinguish e.g. a 503 "still loading" response from a
- * genuine failure.
- */
-export class ApiError extends Error {
-	status: number;
-
-	constructor(message: string, status: number) {
-		super(message);
-		this.name = 'ApiError';
-		this.status = status;
-	}
-}
-
 export interface ApiFetchOptions extends Omit<RequestInit, 'headers'> {
 	/**
 	 * Use auth-only headers (no Content-Type).
@@ -82,7 +67,7 @@ export async function apiFetch<T>(path: string, options: ApiFetchOptions = {}):
 
 	if (!response.ok) {
 		const errorMessage = await parseErrorMessage(response);
-		throw new ApiError(errorMessage, response.status);
+		throw new Error(errorMessage);
 	}
 
 	return response.json() as Promise<T>;
@@ -134,7 +119,7 @@ export async function apiFetchWithParams<T>(
 
 	if (!response.ok) {
 		const errorMessage = await parseErrorMessage(response);
-		throw new ApiError(errorMessage, response.status);
+		throw new Error(errorMessage);
 	}
 
 	return response.json() as Promise<T>;

tools/ui/static/loading.html

@@ -0,0 +1,12 @@
+<!DOCTYPE html>
+<html>
+    <head>
+        <meta http-equiv="refresh" content="5">
+    </head>
+    <body>
+        <div id="loading">
+            The model is loading. Please wait.<br/>
+            The user interface will appear soon.
+        </div>
+    </body>
+</html>