metal : null-check buffer alloc to fix OOM crash (#25371)

* metal : null-check ggml_metal_buffer_init result to avoid OOM crash

ggml_backend_metal_buffer_type_alloc_buffer used the result of
ggml_metal_buffer_init without checking for NULL. ggml_metal_buffer_init
returns NULL when the underlying Metal allocation fails (e.g. an
out-of-memory condition), and the following ggml_metal_buffer_is_shared(res)
call dereferences it, turning a recoverable allocation failure into a hard
crash (EXC_BAD_ACCESS). This is easy to hit on memory-constrained devices
such as iOS when a model/context exceeds the available Metal budget.

Log the failure using the existing GGML_LOG_ERROR convention and return
NULL so the allocator surfaces a diagnosable error up the stack instead of
crashing.

* cont : fix log

---------

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>

.github/workflows/make-release.yml

@@ -84,11 +84,13 @@ jobs:
 
             New version has been released.
 
+            ## Assets
+
             ${{ steps.desc.outputs.nightly }}
 
-            **Web UI:** the `nightly-tag.txt` asset contains the tag of the corresponding nightly release
+            ## More info
 
-            **More info:** [dist : releases and versioning of ggml-org projects](https://github.com/ggml-org/ggml/discussions/1579)
+            - [Releases and versioning of `ggml-org` projects](https://github.com/ggml-org/ggml/discussions/1579)
 
             ## ${{ steps.desc.outputs.changelog_title }}
 

CMakeLists.txt

@@ -4,7 +4,7 @@ include(CheckIncludeFileCXX)
 
 ### llama.cpp version
 set(LLAMA_VERSION_MAJOR 0)
-set(LLAMA_VERSION_MINOR 2)
+set(LLAMA_VERSION_MINOR 3)
 set(LLAMA_VERSION_PATCH 0)
 set(LLAMA_VERSION_BASE "${LLAMA_VERSION_MAJOR}.${LLAMA_VERSION_MINOR}.${LLAMA_VERSION_PATCH}")
 

ggml/CMakeLists.txt

@@ -4,7 +4,7 @@ project("ggml" C CXX ASM)
 
 ### GGML Version
 set(GGML_VERSION_MAJOR 0)
-set(GGML_VERSION_MINOR 21)
+set(GGML_VERSION_MINOR 22)
 set(GGML_VERSION_PATCH 0)
 set(GGML_VERSION_BASE "${GGML_VERSION_MAJOR}.${GGML_VERSION_MINOR}.${GGML_VERSION_PATCH}")
 

ggml/src/ggml-metal/CMakeLists.txt

@@ -11,6 +11,7 @@ ggml_add_backend_library(ggml-metal
                          ggml-metal-common.cpp
                          ggml-metal-context.m
                          ggml-metal-ops.cpp
+                         ggml-metal-tuning.cpp
                         )
 
 target_link_libraries(ggml-metal PRIVATE

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

@@ -1,6 +1,7 @@
 #include "ggml-metal-device.h"
 
 #include "ggml-metal-impl.h"
+#include "ggml-metal-tuning.h"
 
 #include "ggml-impl.h"
 
@@ -1544,6 +1545,8 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
         bool    has_bias,
         bool    has_scap,
         bool    has_kvpad,
+        int32_t nqpsg,
+        int32_t ne,
         int32_t nsg,
         int32_t nwg,
         bool    use_kv_f16,
@@ -1559,11 +1562,17 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
 
     const char * type = use_kv_f16 ? "f16" : ggml_type_name(op->src[1]->type);
 
-    snprintf(base, 256, "kernel_%s_%s_dk%d_dv%d",
+    char qne_suffix[16] = {0};
+    if (!(nqpsg == 1 && ne == ggml_metal_tuning::fa_vec_baseline_ne(dk, dv))) {
+        snprintf(qne_suffix, sizeof(qne_suffix), "_q%d_ne%d", nqpsg, ne);
+    }
+
+    snprintf(base, 256, "kernel_%s_%s_dk%d_dv%d%s",
             "flash_attn_ext_vec",
             type,
             dk,
-            dv);
+            dv,
+            qne_suffix);
 
     snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d",
             base,

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

@@ -207,6 +207,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att
         bool    has_bias,
         bool    has_scap,
         bool    has_kvpad,
+        int32_t nqpsg,
+        int32_t ne,
         int32_t nsg,
         int32_t nwg,
         bool    use_kv_f16,
@@ -257,6 +259,8 @@ enum ggml_metal_device_id {
     GGML_METAL_DEVICE_M5_ULTRA,
 };
 
+const char * ggml_metal_device_id_token(enum ggml_metal_device_id id);
+
 struct ggml_metal_device_props {
     int device;
     int device_phys;
@@ -279,6 +283,7 @@ struct ggml_metal_device_props {
     bool supports_gpu_family_apple7;
 
     enum ggml_metal_device_id device_id;
+    int gpu_family;
 
     int op_offload_min_batch_size;
 };

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

@@ -962,6 +962,34 @@ void ggml_metal_rsets_free(ggml_metal_rsets_t rsets) {
     free(rsets);
 }
 
+static const struct {
+    const char *              name;
+    const char *              token;
+    enum ggml_metal_device_id id;
+} k_metal_devices[] = {
+#define DEV(name, id) { name, #id, id }
+    DEV("M1",       GGML_METAL_DEVICE_M1),
+    DEV("M1 Pro",   GGML_METAL_DEVICE_M1_PRO),
+    DEV("M1 Max",   GGML_METAL_DEVICE_M1_MAX),
+    DEV("M1 Ultra", GGML_METAL_DEVICE_M1_ULTRA),
+    DEV("M2",       GGML_METAL_DEVICE_M2),
+    DEV("M2 Pro",   GGML_METAL_DEVICE_M2_PRO),
+    DEV("M2 Max",   GGML_METAL_DEVICE_M2_MAX),
+    DEV("M2 Ultra", GGML_METAL_DEVICE_M2_ULTRA),
+    DEV("M3",       GGML_METAL_DEVICE_M3),
+    DEV("M3 Pro",   GGML_METAL_DEVICE_M3_PRO),
+    DEV("M3 Max",   GGML_METAL_DEVICE_M3_MAX),
+    DEV("M3 Ultra", GGML_METAL_DEVICE_M3_ULTRA),
+    DEV("M4",       GGML_METAL_DEVICE_M4),
+    DEV("M4 Pro",   GGML_METAL_DEVICE_M4_PRO),
+    DEV("M4 Max",   GGML_METAL_DEVICE_M4_MAX),
+    DEV("M5",       GGML_METAL_DEVICE_M5),
+    DEV("M5 Pro",   GGML_METAL_DEVICE_M5_PRO),
+    DEV("M5 Max",   GGML_METAL_DEVICE_M5_MAX),
+    DEV("M5 Ultra", GGML_METAL_DEVICE_M5_ULTRA),
+#undef DEV
+};
+
 static enum ggml_metal_device_id ggml_metal_device_id_parse(const char * name) {
     if (!name) {
         return GGML_METAL_DEVICE_GENERIC;
@@ -973,39 +1001,23 @@ static enum ggml_metal_device_id ggml_metal_device_id_parse(const char * name) {
     }
     const char * suffix = name + sizeof(prefix) - 1;
 
-    static const struct {
-        const char * name;
-        enum ggml_metal_device_id id;
-    } table[] = {
-        {"M1",       GGML_METAL_DEVICE_M1},
-        {"M1 Pro",   GGML_METAL_DEVICE_M1_PRO},
-        {"M1 Max",   GGML_METAL_DEVICE_M1_MAX},
-        {"M1 Ultra", GGML_METAL_DEVICE_M1_ULTRA},
-        {"M2",       GGML_METAL_DEVICE_M2},
-        {"M2 Pro",   GGML_METAL_DEVICE_M2_PRO},
-        {"M2 Max",   GGML_METAL_DEVICE_M2_MAX},
-        {"M2 Ultra", GGML_METAL_DEVICE_M2_ULTRA},
-        {"M3",       GGML_METAL_DEVICE_M3},
-        {"M3 Pro",   GGML_METAL_DEVICE_M3_PRO},
-        {"M3 Max",   GGML_METAL_DEVICE_M3_MAX},
-        {"M3 Ultra", GGML_METAL_DEVICE_M3_ULTRA},
-        {"M4",       GGML_METAL_DEVICE_M4},
-        {"M4 Pro",   GGML_METAL_DEVICE_M4_PRO},
-        {"M4 Max",   GGML_METAL_DEVICE_M4_MAX},
-        {"M5",       GGML_METAL_DEVICE_M5},
-        {"M5 Pro",   GGML_METAL_DEVICE_M5_PRO},
-        {"M5 Max",   GGML_METAL_DEVICE_M5_MAX},
-        {"M5 Ultra", GGML_METAL_DEVICE_M5_ULTRA},
-    };
-
-    for (size_t i = 0; i < sizeof(table)/sizeof(table[0]); ++i) {
-        if (strcmp(suffix, table[i].name) == 0) {
-            return table[i].id;
+    for (size_t i = 0; i < sizeof(k_metal_devices)/sizeof(k_metal_devices[0]); ++i) {
+        if (strcmp(suffix, k_metal_devices[i].name) == 0) {
+            return k_metal_devices[i].id;
         }
     }
     return GGML_METAL_DEVICE_GENERIC;
 }
 
+const char * ggml_metal_device_id_token(enum ggml_metal_device_id id) {
+    for (size_t i = 0; i < sizeof(k_metal_devices)/sizeof(k_metal_devices[0]); ++i) {
+        if (k_metal_devices[i].id == id) {
+            return k_metal_devices[i].token;
+        }
+    }
+    return "GGML_METAL_DEVICE_GENERIC";
+}
+
 ggml_metal_device_t ggml_metal_device_init(int device, int n_devices) {
     ggml_metal_device_t dev = calloc(1, sizeof(struct ggml_metal_device));
 
@@ -1220,7 +1232,8 @@ ggml_metal_device_t ggml_metal_device_init(int device, int n_devices) {
             {
                 for (int i = MTLGPUFamilyApple1 + 20; i >= MTLGPUFamilyApple1; --i) {
                     if ([dev->mtl_device supportsFamily:i]) {
-                        GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyApple%d  (%d)\n", __func__, i - (int) MTLGPUFamilyApple1 + 1, i);
+                        dev->props.gpu_family = i - (int) MTLGPUFamilyApple1 + 1;
+                        GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyApple%d  (%d)\n", __func__, dev->props.gpu_family, i);
                         break;
                     }
                 }

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

@@ -7,6 +7,7 @@
 #include "ggml-metal-impl.h"
 #include "ggml-metal-common.h"
 #include "ggml-metal-device.h"
+#include "ggml-metal-tuning.h"
 
 #include <cassert>
 #include <algorithm>
@@ -3346,12 +3347,18 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
 #undef FATTN_SMEM
     } else {
         // half4x4 kernel
-        const int nqptg = OP_FLASH_ATTN_EXT_VEC_NQPSG; // queries per threadgroup
+        auto cfg = ggml_metal_tuning::fa_vec_pick(
+                props_dev->device_id,
+                props_dev->gpu_family,
+                (int) op->src[1]->type,
+                (int) ne00, (int) ne20,   // dk, dv (ne00 == dk for FA)
+                ne11, ne01);
+        int nqptg = cfg.Q;                             // queries per threadgroup
         const int ncpsg = OP_FLASH_ATTN_EXT_VEC_NCPSG; // cache values per simdgroup !! sync with kernel template arguments !!
         const int nhptg = 1;                           // heads per threadgroup
 
         GGML_ASSERT(nqptg <= 32);
-        GGML_ASSERT(nqptg  % 1  == 0);
+        GGML_ASSERT(nqptg == 1 || nqptg == 2 || nqptg == 4);  // only instantiated Q values
         GGML_ASSERT(ncpsg  % 32 == 0);
 
         bool need_sync = false;
@@ -3410,7 +3417,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
         // ne20*(nsg)
         // each simdgroup has a full f32 head vector in shared mem to accumulate results
         //
-#define FATTN_SMEM(nsg) (GGML_PAD(((GGML_PAD(ne00, 128) + 4*ncpsg + 2*GGML_PAD(ne20, 128))*(nsg))*(sizeof(float)/2), 16))
+#define FATTN_SMEM(nsg) (GGML_PAD(((GGML_PAD(ne00, 128) + 4*ncpsg + 2*GGML_PAD(ne20, 128))*(nsg)*nqptg)*(sizeof(float)/2), 16))
 
         int64_t nsg = 1;
 
@@ -3430,6 +3437,12 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
             }
         }
 
+        // fall back to baseline (Q=1) if the tuned config exceeds threadgroup memory
+        if ((size_t) FATTN_SMEM(nsg) > props_dev->max_theadgroup_memory_size) {
+            cfg   = ggml_metal_tuning::fa_vec_baseline_cfg((int) ne00, (int) ne20);
+            nqptg = cfg.Q;  // = 1
+        }
+
         const int32_t ns10 = nb11_attn/nb10_attn;
         const int32_t ns20 = nb21_attn/nb20_attn;
 
@@ -3468,7 +3481,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
             /*.logit_softcap =*/ logit_softcap,
         };
 
-        auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, nwg, use_kv_f16, ns10, ns20);
+        auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nqptg, cfg.NE, nsg, nwg, use_kv_f16, ns10, ns20);
 
         GGML_ASSERT(nsg*32 <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
 

ggml/src/ggml-metal/ggml-metal-tuning.h

@@ -0,0 +1,77 @@
+#pragma once
+
+#include "ggml-metal-device.h"  // enum ggml_metal_device_id
+#include "ggml.h"
+
+#include <cstdint>
+#include <vector>
+
+namespace ggml_metal_tuning {
+
+// FA vec selection buckets. ne01 (query rows) splits decode (==1) from batch (>=2), the
+// batch side refined into {2,3,4,5}: Q>1 reuses one K/V load across rows, so it only pays
+// off once ne01 aligns with Q. ne11 (KV length) is bucketed too, as the Q>1 crossover is
+// head-size dependent (small dk crosses late, large dk wins even at short KV).
+constexpr int FA_VEC_NE11_BUCKETS[] = { 1024, 4096, 16384 };
+constexpr int FA_VEC_NE01_BUCKETS[] = { 2, 3, 4, 5 };
+
+int fa_vec_ne11_bucket(int64_t ne11);
+int fa_vec_ne01_bucket(int64_t ne01);
+
+// NE baked into each (dk,dv) baseline instantiation in kernels/fa.metal.
+// Hand-maintained mirror; keep in sync with those instantiations.
+// The Metal test slice covers every legal config for dk=128 and dk=576.
+int fa_vec_baseline_ne(int dk, int dv);
+
+// Tuned table has two row kinds. Exact rows key a (ne11_b, ne01_b) bucket. Default rows
+// collapse ne11 over one ne01 domain: ne11_b == FA_VEC_NE11_DEFAULT and ne01_b holds the
+// domain. fa_vec_pick tries exact bucket -> domain default -> baseline; short KV
+// (ne11 < FA_VEC_NE11_BUCKETS[0]) always uses baseline.
+constexpr int8_t FA_VEC_NE11_DEFAULT  = -1;
+constexpr int8_t FA_VEC_DOMAIN_DECODE = 0;  // ne01 == 1
+constexpr int8_t FA_VEC_DOMAIN_BATCH  = 1;  // ne01 >= 2
+
+struct fa_vec_key_t {
+    int8_t  device_id;
+    int8_t  dtype;
+    int16_t dk;
+    int16_t dv;
+    int8_t  ne11_b;
+    int8_t  ne01_b;
+};
+
+static_assert(sizeof(fa_vec_key_t) == 8, "fa_vec_key_t must be tightly packed for memcmp");
+
+struct fa_vec_cfg_t {
+    int8_t Q;
+    int8_t NE;
+};
+
+struct fa_vec_entry_t {
+    fa_vec_key_t key;
+    fa_vec_cfg_t cfg;
+};
+
+// legal NE values for a (dk,dv): NL = 32/NE, require (dk/4)%NL==0 && (dv/4)%NL==0.
+// single source shared by the offline tuner and test-backend-ops.
+inline std::vector<int> fa_vec_legal_ne(int dk, int dv) {
+    std::vector<int> r;
+    for (int ne : { 1, 2, 4 }) {
+        const int nl = 32 / ne;
+        if ((dk / 4) % nl == 0 && (dv / 4) % nl == 0) {
+            r.push_back(ne);
+        }
+    }
+    return r;
+}
+
+// test/tune-only override; when set, fa_vec_pick returns it directly.
+void         fa_vec_set_override(fa_vec_cfg_t cfg);
+void         fa_vec_clear_override();
+fa_vec_cfg_t fa_vec_baseline_cfg(int dk, int dv);
+
+// device_id selects a per-SKU row; on a miss, gpu_family (0 if unknown) maps to a representative
+// SKU and the table is retried. No match -> baseline.
+fa_vec_cfg_t fa_vec_pick(enum ggml_metal_device_id device_id, int gpu_family, int dtype, int dk, int dv, int64_t ne11, int64_t ne01);
+
+}  // namespace ggml_metal_tuning

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

@@ -6,6 +6,7 @@
 #include "ggml-metal-device.h"
 #include "ggml-metal-context.h"
 #include "ggml-metal-ops.h"
+#include "ggml-metal-tuning.h"
 
 #include <mutex>
 #include <string>
@@ -203,6 +204,11 @@ static ggml_backend_buffer_t ggml_backend_metal_buffer_type_alloc_buffer(ggml_ba
     ggml_metal_device_t ctx_dev = (ggml_metal_device_t)buft->device->context;
     ggml_metal_buffer_t res = ggml_metal_buffer_init(ctx_dev, size, shared);
 
+    if (res == NULL) {
+        GGML_LOG_ERROR("%s: failed to allocate Metal buffer of %zu bytes (out of memory)\n", __func__, size);
+        return NULL;
+    }
+
     ggml_backend_buffer_i buf_i = ggml_metal_buffer_is_shared(res)
         ? ggml_backend_metal_buffer_shared_i
         : ggml_backend_metal_buffer_private_i;
@@ -870,10 +876,55 @@ static ggml_backend_feature * ggml_backend_metal_get_features(ggml_backend_reg_t
     GGML_UNUSED(reg);
 }
 
+// test/tune-only override for the FA vec (Q, NE) selection, reached via proc_address.
+static void ggml_backend_metal_tuning_set_fa_vec_override(int Q, int NE) {
+    ggml_metal_tuning::fa_vec_set_override({ (int8_t) Q, (int8_t) NE });
+}
+
+static void ggml_backend_metal_tuning_clear_fa_vec_override(void) {
+    ggml_metal_tuning::fa_vec_clear_override();
+}
+
+static int ggml_backend_metal_tuning_fa_vec_ne11_bucket(int64_t ne11) {
+    return ggml_metal_tuning::fa_vec_ne11_bucket(ne11);
+}
+
+static int ggml_backend_metal_tuning_fa_vec_ne01_bucket(int64_t ne01) {
+    return ggml_metal_tuning::fa_vec_ne01_bucket(ne01);
+}
+
+static int ggml_backend_metal_tuning_fa_vec_baseline_ne(int dk, int dv) {
+    return ggml_metal_tuning::fa_vec_baseline_ne(dk, dv);
+}
+
+static const char * ggml_backend_metal_tuning_device_token(ggml_backend_dev_t dev) {
+    ggml_metal_device_t ctx_dev = (ggml_metal_device_t)dev->context;
+
+    return ggml_metal_device_id_token(ggml_metal_device_get_props(ctx_dev)->device_id);
+}
+
 static void * ggml_backend_metal_get_proc_address(ggml_backend_reg_t reg, const char * name) {
     if (strcmp(name, "ggml_backend_get_features") == 0) {
         return (void *)ggml_backend_metal_get_features;
     }
+    if (strcmp(name, "ggml_backend_metal_tuning_set_fa_vec_override") == 0) {
+        return (void *)ggml_backend_metal_tuning_set_fa_vec_override;
+    }
+    if (strcmp(name, "ggml_backend_metal_tuning_clear_fa_vec_override") == 0) {
+        return (void *)ggml_backend_metal_tuning_clear_fa_vec_override;
+    }
+    if (strcmp(name, "ggml_backend_metal_tuning_fa_vec_ne11_bucket") == 0) {
+        return (void *)ggml_backend_metal_tuning_fa_vec_ne11_bucket;
+    }
+    if (strcmp(name, "ggml_backend_metal_tuning_fa_vec_ne01_bucket") == 0) {
+        return (void *)ggml_backend_metal_tuning_fa_vec_ne01_bucket;
+    }
+    if (strcmp(name, "ggml_backend_metal_tuning_fa_vec_baseline_ne") == 0) {
+        return (void *)ggml_backend_metal_tuning_fa_vec_baseline_ne;
+    }
+    if (strcmp(name, "ggml_backend_metal_tuning_device_token") == 0) {
+        return (void *)ggml_backend_metal_tuning_device_token;
+    }
 
     return NULL;
 

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

@@ -6018,6 +6018,11 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
                     a->ne[0] > 128 && a->ne[2] == 1 && src0_type == GGML_TYPE_F16) {
                     return false;
                 }
+
+                if (src0_type == GGML_TYPE_TQ2_0) {
+                    return false;
+                }
+
                 return true;
             }
         case GGML_OP_OUT_PROD:
@@ -6068,6 +6073,9 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
 
         case GGML_OP_SET_ROWS:
             {
+                if (op->type == GGML_TYPE_TQ2_0) {
+                    return false;
+                }
                 auto res = (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16 ||
                             op->src[0]->type == GGML_TYPE_BF16) &&
                            (op->src[1]->type == GGML_TYPE_I64 || op->src[1]->type == GGML_TYPE_I32);
@@ -6186,11 +6194,16 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
                         src1_type == GGML_TYPE_IQ3_XXS ||
                         src1_type == GGML_TYPE_IQ3_S ||
                         src1_type == GGML_TYPE_IQ1_S ||
-                        src1_type == GGML_TYPE_IQ1_M) {
+                        src1_type == GGML_TYPE_IQ1_M ||
+                        src1_type == GGML_TYPE_TQ2_0) {
                         return false;
                     }
                 }
 
+                if (src0_type == GGML_TYPE_TQ2_0 || src1_type == GGML_TYPE_TQ2_0) {
+                    return false;
+                }
+
                 return true;
             }
         case GGML_OP_REPEAT_BACK:

ggml/src/ggml-sycl/mmvq.cpp

@@ -2559,7 +2559,7 @@ void ggml_sycl_op_mul_mat_vec_q(ggml_backend_sycl_context & ctx, const ggml_tens
                 }
                 break;
             default:
-                GGML_ABORT("fatal error: unsupport data type=%s\n", ggml_type_name(src0->type));
+                GGML_ABORT("fatal error: unsupport src0 data type %s\n", ggml_type_name(src0->type));
         }
     }
     GGML_UNUSED(src1);

ggml/src/ggml-sycl/set_rows.cpp

@@ -546,7 +546,8 @@ static void set_rows_sycl(ggml_backend_sycl_context & ctx, const ggml_tensor * s
                 stream);
             break;
         default:
-            GGML_ABORT("Unsupported tensor type!");
+            GGML_ABORT("Unsupported tensor type: src0 %s src1 %s dst %s", ggml_type_name(dst->src[0]->type),
+                ggml_type_name(dst->src[1]->type), ggml_type_name(dst->type));
             break;
     }
 }

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

@@ -34,11 +34,9 @@ var<uniform> params: Params;
 var<workgroup> shmem_idx: array<u32, WG_SIZE>;
 
 #if ORDER == 0
-#define EXTREME_VALUE 1e30
 #define SWAP_COMPARE_UP >
 #define SWAP_COMPARE_DOWN <
 #else
-#define EXTREME_VALUE -1e30
 #define SWAP_COMPARE_UP <
 #define SWAP_COMPARE_DOWN >
 #endif
@@ -78,11 +76,9 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
                 let dir_up = (lid.x & k) == 0;
                 let a_idx = shmem_idx[lid.x];
                 let b_idx = shmem_idx[ixj];
-                let a_val = select(EXTREME_VALUE, src[row_base + a_idx], a_idx < params.src_ne0);
-                let b_val = select(EXTREME_VALUE, src[row_base + b_idx], b_idx < params.src_ne0);
                 let should_swap = select(
-                    (a_val SWAP_COMPARE_DOWN b_val),
-                    (a_val SWAP_COMPARE_UP b_val),
+                    b_idx >= params.src_ne0 || (a_idx < params.src_ne0 && src[row_base + a_idx] SWAP_COMPARE_DOWN src[row_base + b_idx]),
+                    a_idx >= params.src_ne0 || (b_idx < params.src_ne0 && src[row_base + a_idx] SWAP_COMPARE_UP src[row_base + b_idx]),
                     dir_up);
                 if (should_swap) {
                     shmem_idx[lid.x] = b_idx;

scripts/make-release-summary.txt

@@ -0,0 +1,47 @@
+Take a look at the changelog between the current version and the previous version - use the `./scripts/make-release-desc.sh [current-version]` to obtain it.
+
+Write a summary of the change log in a few sections:
+
+```
+## Overview
+
+[an overview using 1 to 3 sentences (no line breaks)]
+
+### API changes (if applicable)
+
+[summarize any API changes to `/include/*`, `/tools/mtmd/mtmd.h` and `/tools/server`]
+
+### New models (if applicable)
+
+[summarize new models added to the `src/models/` directory]
+
+### Core changes (if applicable)
+
+[summarize the changes to `/src/*`
+
+### Multi-modality changes (if applicable)
+
+[summarize the changes to `/tools/mtmd/`]
+
+### Server changes (if applicable)
+
+[summarize the changes to `/tools/server/`]
+
+### UI changes (if applicable)
+
+[summarize the changes to `/tools/ui/`]
+
+### ggml changes (if applicable)
+
+[if the version has been updated, link to the respective `ggml` releases on Github, f.ex `https://github.com/ggml-org/ggml/releases/tag/v0.22.0`. for each version bump, lookup the release description and copy the summary here]
+
+```
+
+Guidelines:
+
+- All bullet point in the summary should be concise and rarely exceed a single line of 120 characters
+- Avoid repeating `ggml`-specific changes - these should already be covered by the `ggml` release links
+- Provide PR link for each bullet point where possible
+- Don't add bullet point to state that there are no API changes in some module
+
+Output just the summary in a markdown block, without any extra text.

scripts/sync-ggml.last

@@ -1 +1 @@
-8599e0ea3756c4bac4ef813af2241cb1a8bbfb0b
+36da57138425487184aa1da2eee2cde155909c6f

src/llama-grammar.cpp

@@ -172,6 +172,7 @@ static std::pair<uint32_t, const char *> parse_char(const char * src) {
             case '"':
             case '[':
             case ']':
+            case '-':
                       return std::make_pair(src[1], src + 2);
             default:
                       throw std::runtime_error(std::string("unknown escape at ") + src);

tests/test-backend-ops.cpp

@@ -10573,6 +10573,101 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_from_file(const c
     return test_cases;
 }
 
+// ---- FA vec (Q,NE): forced-config numerical slice (Metal only) ----
+using set_fa_vec_override_t   = void (*)(int, int);
+using clear_fa_vec_override_t = void (*)(void);
+
+// NL = 32/NE must divide both dk/4 and dv/4.
+static std::vector<int> fa_vec_legal_ne(int dk, int dv) {
+    std::vector<int> r;
+    for (int ne : {1, 2, 4}) {
+        const int nl = 32 / ne;
+        if ((dk/4) % nl == 0 && (dv/4) % nl == 0) {
+            r.push_back(ne);
+        }
+    }
+    return r;
+}
+
+static bool op_names_filter_selects(const char * op_names_filter, const char * op_name) {
+    if (!op_names_filter) {
+        return true;
+    }
+    std::string_view filter(op_names_filter);
+    while (!filter.empty()) {
+        auto comma_pos = filter.find_first_of(',');
+        const auto lparen_pos = filter.find_first_of('(');
+        std::string_view entry;
+        if (lparen_pos < comma_pos) {
+            const auto rparen_pos = filter.find_first_of(')');
+            comma_pos = filter.find_first_of(',', rparen_pos);
+            entry = filter.substr(0, lparen_pos);
+        } else {
+            entry = filter.substr(0, comma_pos);
+        }
+        if (entry == op_name) {
+            return true;
+        }
+        filter = comma_pos != std::string_view::npos ? filter.substr(comma_pos + 1) : "";
+    }
+    return false;
+}
+
+// Covers padded rows, sinks, kvpad, multi-SIMDgroup reduction, quantized K/V, and MLA views.
+// The override is backend-global, so this runs after all parallel workers have joined.
+static bool run_fa_vec_slice(ggml_backend_t backend, ggml_backend_t backend_cpu, const char * op_names_filter) {
+    if (!op_names_filter_selects(op_names_filter, "FLASH_ATTN_EXT")) {
+        return true;
+    }
+
+    auto * reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
+
+    auto set_ov   = (set_fa_vec_override_t)   ggml_backend_reg_get_proc_address(reg, "ggml_backend_metal_tuning_set_fa_vec_override");
+    auto clear_ov = (clear_fa_vec_override_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_metal_tuning_clear_fa_vec_override");
+    if (!set_ov || !clear_ov) {
+        return true;  // not the Metal backend: nothing to force
+    }
+
+    struct shape_t { int dk, dv; };
+    const shape_t   shapes[] = { { 128, 128 }, { 576, 512 } };  // mainstream head size + MLA shared K/V view
+    const int       ne01_pts[] = { 1, 3 };                      // decode, and padded rows for Q=2 and Q=4
+    const int       ne11_pts[] = { 512, 4097 };                 // nsg=1, and nsg>=2 together with kvpad
+    const ggml_type types[]    = { GGML_TYPE_F16, GGML_TYPE_Q4_0 };
+
+    int n_run = 0, n_fail = 0;
+    for (auto s : shapes) {
+        for (int ne : fa_vec_legal_ne(s.dk, s.dv)) {
+            for (int Q : { 1, 2, 4 }) {
+                for (ggml_type type_kv : types) {
+                    for (bool sinks : { false, true }) {
+                        for (int ne01 : ne01_pts) {
+                            for (int ne11 : ne11_pts) {
+                                set_ov(Q, ne);
+                                test_flash_attn_ext tc(s.dk, s.dv, /*nh=*/4, { 1, 1 }, /*kv=*/ne11, /*nb=*/ne01,
+                                                       /*mask=*/true, sinks, 0.0f, 0.0f, GGML_PREC_F32,
+                                                       type_kv, type_kv);
+                                auto st = tc.eval(backend, backend_cpu, "FLASH_ATTN_EXT", nullptr);
+                                clear_ov();
+
+                                if (st == test_status_t::FAIL) {
+                                    printf("  FAIL fa_vec slice: dk=%d dv=%d Q=%d ne=%d type=%s ne01=%d ne11=%d sinks=%d\n",
+                                           s.dk, s.dv, Q, ne, ggml_type_name(type_kv), ne01, ne11, (int) sinks);
+                                    n_fail++;
+                                }
+                                n_run++;
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    printf("  fa_vec (Q,NE) slice: %d cases run, %d failed\n", n_run, n_fail);
+
+    return n_fail == 0;
+}
+
 static bool test_backend(ggml_backend_t backend, ggml_backend_dev_t dev, test_mode mode, const char * op_names_filter, const char * params_filter,
                          printer * output_printer, const char * test_file_path, int parallel_workers) {
     auto filter_test_cases = [](std::vector<std::unique_ptr<test_case>> & test_cases, const char * params_filter) {
@@ -10710,7 +10805,9 @@ static bool test_backend(ggml_backend_t backend, ggml_backend_dev_t dev, test_mo
         output_printer->print_summary(test_summary_info(n_ok, tests_run, false));
         output_printer->print_failed_tests(failed_tests);
 
-        return n_ok == tests_run;
+        const bool slice_ok = run_fa_vec_slice(backend, backend_cpu.get(), op_names_filter);
+
+        return n_ok == tests_run && slice_ok;
     }
 
     if (mode == MODE_GRAD) {

tools/CMakeLists.txt

@@ -38,5 +38,8 @@ else()
         add_subdirectory(export-lora)
     endif()
     add_subdirectory(fit-params)
+    if (GGML_METAL)
+        add_subdirectory(tuning)
+    endif()
     add_subdirectory(results)
 endif()

tools/tuning/CMakeLists.txt

@@ -0,0 +1,10 @@
+set(TARGET ggml-metal-tuning)
+
+add_executable(${TARGET} main.cpp bench.cpp fa-vec.cpp)
+target_link_libraries(${TARGET} PRIVATE ggml ${CMAKE_THREAD_LIBS_INIT})
+target_compile_features(${TARGET} PRIVATE cxx_std_17)
+target_include_directories(${TARGET} PRIVATE ${CMAKE_SOURCE_DIR}/ggml/src/ggml-metal)
+
+if(LLAMA_TOOLS_INSTALL)
+    install(TARGETS ${TARGET} RUNTIME)
+endif()

tools/tuning/README.md

@@ -0,0 +1,63 @@
+# ggml-metal-tuning
+
+Offline kernel tuner for the Metal backend.
+It sweeps a kernel's config grid on the machine it runs on and prints pasteable table rows for `ggml/src/ggml-metal/ggml-metal-tuning.cpp`.
+
+This is not a test: it never reports pass/fail on performance.
+A non-zero exit code means bad arguments or a wrong environment (no Metal device, missing proc bridges), never a perf result.
+
+| tuner | tunes | table |
+|---|---|---|
+| `fa-vec` | flash-attn vec `(Q, NE)` per `(dtype, head size, KV depth, batch width)` | `fa_vec_tuned_table` |
+
+## Adding a device to the FA-vec table
+
+Build on the target machine:
+
+```bash
+cmake -B build -DGGML_METAL=ON
+cmake --build build --target ggml-metal-tuning -j
+cmake --build build --target test-backend-ops -j
+```
+
+Sweep the grid (6 dtypes x 10 head sizes x 4 KV depths x 9 batch widths; a few hours):
+
+```bash
+./build/bin/ggml-metal-tuning fa-vec > fa_vec_rows.txt 2> fa_vec_sweep.log
+```
+
+`fa_vec_rows.txt` holds nothing but table rows, ready to paste into `fa_vec_tuned_table`: the min-max-regret target, the aggregate benefit gate, the short-KV drop and the pointwise compression are already applied.
+A config represents a bucket only if it is no slower than the baseline config at every point that bucket covers, so a config that wins on average but loses at one batch width leaves its bucket at baseline.
+`fa_vec_sweep.log` holds the per-cell timings, bucket coverage, noise floor, any cooldown activity, and every config the no-harm rule refused together with the point that refused it.
+Post both: the log is what makes the rows reviewable.
+
+Long sweeps can be split.
+`--dtype f16,q4_0` and `--dk 128,192` restrict the grid, and the emitted rows for one `(dtype, head size)` do not depend on the others.
+Concatenating the shard outputs in the order the full grid would visit them gives the same rows a single run prints.
+
+Then validate the numerics, where Metal is compared against the CPU reference:
+
+```bash
+./build/bin/test-backend-ops test -o FLASH_ATTN_EXT -b MTL0
+```
+
+This forces every legal `(Q, NE)` on `dk=128` and `dk=576`.
+The tuner itself does no numerical checks, so the other head sizes have no automated numerical coverage.
+
+If the device is not in `enum ggml_metal_device_id` yet, register it in `ggml/src/ggml-metal/ggml-metal-device.{h,m}` first.
+The tuner emits whatever token the runtime reports for the machine, so an unregistered device emits `GGML_METAL_DEVICE_GENERIC` and its rows would apply to every unknown device.
+
+## Thermal throttling
+
+Long sweeps heat the GPU, and a throttled measurement is indistinguishable from a slow kernel.
+The tuner re-measures a fixed baseline config every four candidates as an anchor.
+When the anchor drifts more than `--cool-drift` (10% by default) from the coolest anchor seen in that cell, the tuner:
+
+1. discards every candidate measured since the last clean anchor,
+2. sleeps with exponential backoff until the anchor comes back within `--cool-eps` (3%),
+3. re-measures the discarded candidates.
+
+If it cannot cool down within `--cool-max-wait` seconds, or a cell needs more than `--cool-max-retry` rounds, that cell is dropped from the table and reported on stderr.
+
+`--no-cooldown` only warns on drift and keeps the measurement.
+Use it to reproduce a sweep taken without cooling.

tools/tuning/bench.cpp

@@ -0,0 +1,234 @@
+#include "bench.h"
+
+#include <algorithm>
+#include <chrono>
+#include <cmath>
+#include <cstdio>
+#include <thread>
+#include <utility>
+
+perf_cell build_perf_cell(ggml_backend_t          backend,
+                          const build_graph_fn &  build,
+                          const init_tensors_fn & init,
+                          const op_flops_fn &     flops) {
+    perf_cell cell;
+
+    const size_t graph_nodes = 1024;
+
+    ggml_init_params params = {
+        /* .mem_size  = */ ggml_tensor_overhead() * 128 + ggml_graph_overhead_custom(graph_nodes, false),
+        /* .mem_base  = */ NULL,
+        /* .no_alloc  = */ true,
+    };
+
+    cell.ctx.reset(ggml_init(params));
+    GGML_ASSERT(cell.ctx);
+
+    ggml_tensor * out = build(cell.ctx.get());
+    if (!out || !ggml_backend_supports_op(backend, out)) {
+        return cell;
+    }
+
+    cell.buf.reset(ggml_backend_alloc_ctx_tensors(cell.ctx.get(), backend));
+    if (!cell.buf) {
+        return cell;
+    }
+
+    init(cell.ctx.get());
+
+    cell.gf = ggml_new_graph_custom(cell.ctx.get(), graph_nodes, false);
+    ggml_build_forward_expand(cell.gf, out);
+
+    // replicate the op to amortize overhead (target ~50 GFLOP/compute, capped to bound graph size)
+    cell.n_runs            = 1;
+    const uint64_t n_flops = flops(out);
+    if (n_flops > 0) {
+        const uint64_t target_flops = 50ULL * 1000 * 1000 * 1000;
+        const int      cap          = 512;
+        const int      by_flops     = (int) std::min<int64_t>(cap, (int64_t) (target_flops / n_flops));
+        cell.n_runs =
+            std::max(1, std::min<int>(by_flops, (int) (ggml_graph_size(cell.gf) - ggml_graph_n_nodes(cell.gf))));
+    }
+    for (int i = 1; i < cell.n_runs; ++i) {
+        ggml_graph_add_node(cell.gf, out);
+    }
+
+    return cell;
+}
+
+double time_cell_median(ggml_backend_t backend, const perf_cell & cell, int reps) {
+    if (cell.gf == nullptr) {
+        return -1.0;
+    }
+
+    ggml_backend_graph_compute(backend, cell.gf);  // warmup (compiles the pipeline for this config)
+    ggml_backend_synchronize(backend);
+
+    std::vector<double> samples;
+    samples.reserve(reps);
+    for (int r = 0; r < reps; ++r) {
+        const int64_t t0 = ggml_time_us();
+        ggml_backend_graph_compute(backend, cell.gf);
+        ggml_backend_synchronize(backend);
+        samples.push_back((double) (ggml_time_us() - t0));
+    }
+    std::nth_element(samples.begin(), samples.begin() + samples.size() / 2, samples.end());
+
+    return samples[samples.size() / 2] / cell.n_runs;
+}
+
+static double measure_one(ggml_backend_t             backend,
+                          const perf_cell &          cell,
+                          int                        reps,
+                          const set_candidate_fn &   set_cand,
+                          const clear_candidate_fn & clear_cand,
+                          int                        cand) {
+    set_cand(cand);
+    const double t = time_cell_median(backend, cell, reps);
+    clear_cand();
+
+    return t;
+}
+
+// waits for the anchor to come back within eps of anchor_ref, with exponential backoff.
+// returns the converged anchor, or -1 if it never converged within max_wait.
+static double cool_until_steady(ggml_backend_t             backend,
+                                const perf_cell &          cell,
+                                int                        reps,
+                                const set_candidate_fn &   set_cand,
+                                const clear_candidate_fn & clear_cand,
+                                int                        baseline_cand,
+                                double &                   anchor_ref,
+                                const cooldown_opts &      cool,
+                                const char *               cell_label) {
+    int total_wait = 0;
+
+    for (int sleep_s = 2; total_wait < cool.max_wait; sleep_s = std::min(sleep_s * 2, 32)) {
+        const int this_wait = std::min(sleep_s, cool.max_wait - total_wait);
+
+        fprintf(stderr, "# COOL sleeping %ds (%ds/%ds) %s\n", this_wait, total_wait + this_wait, cool.max_wait,
+                cell_label);
+        std::this_thread::sleep_for(std::chrono::seconds(this_wait));
+        total_wait += this_wait;
+
+        const double a = measure_one(backend, cell, reps, set_cand, clear_cand, baseline_cand);
+        if (a <= 0.0) {
+            continue;
+        }
+
+        // a faster anchor means the machine got cooler than anything seen so far: adopt it
+        if (a < anchor_ref) {
+            anchor_ref = a;
+        }
+
+        if (a <= anchor_ref * (1.0 + cool.eps)) {
+            fprintf(stderr, "# COOL steady after %ds %s\n", total_wait, cell_label);
+            return a;
+        }
+    }
+
+    fprintf(stderr, "# COOL gave up after %ds %s\n", total_wait, cell_label);
+
+    return -1.0;
+}
+
+cell_result measure_cell(ggml_backend_t             backend,
+                         const perf_cell &          cell,
+                         int                        reps,
+                         const std::vector<int> &   order,
+                         const set_candidate_fn &   set_cand,
+                         const clear_candidate_fn & clear_cand,
+                         int                        baseline_cand,
+                         const cooldown_opts &      cool,
+                         const char *               cell_label) {
+    cell_result res;
+    res.t.assign(order.size(), 0.0);
+
+    double anchor_ref = 0.0;
+
+    // anchors accepted as clean, as (value, position in order[]). the dirty window starts
+    // at the position of the last anchor still within eps of anchor_ref, so a downward
+    // drift (anchor_ref dropping) naturally widens the window to the whole cell.
+    std::vector<std::pair<double, size_t>> anchors;
+
+    auto window_start = [&]() -> size_t {
+        for (size_t i = anchors.size(); i-- > 0;) {
+            if (anchors[i].first <= anchor_ref * (1.0 + cool.eps)) {
+                return anchors[i].second;
+            }
+        }
+        return 0;  // no clean anchor left -> the whole cell is suspect
+    };
+
+    int retries_left = cool.max_retry;
+
+    for (size_t i = 0; i < order.size(); ++i) {
+        res.t[order[i]] = measure_one(backend, cell, reps, set_cand, clear_cand, order[i]);
+
+        if (i % 4 != 0) {
+            continue;
+        }
+
+        const double a = measure_one(backend, cell, reps, set_cand, clear_cand, baseline_cand);
+        if (a <= 0.0) {
+            continue;
+        }
+
+        res.anchor_min = res.anchor_min > 0.0 ? std::min(res.anchor_min, a) : a;
+        res.anchor_max = std::max(res.anchor_max, a);
+
+        if (anchor_ref == 0.0) {
+            anchor_ref = a;
+            anchors.push_back({ a, i });
+            continue;
+        }
+
+        const double drift = std::fabs(a - anchor_ref) / anchor_ref;
+
+        // a cooler anchor than any so far becomes the reference: whatever was measured
+        // before it was measured on a hotter machine
+        if (a < anchor_ref) {
+            anchor_ref = a;
+        }
+
+        if (drift <= cool.drift) {
+            anchors.push_back({ a, i });
+            continue;
+        }
+
+        fprintf(stderr, "# WARN throttling? anchor drift %.1f%% %s\n", 100.0 * drift, cell_label);
+
+        if (!cool.enabled) {
+            anchors.push_back({ a, i });
+            continue;
+        }
+
+        if (retries_left <= 0) {
+            fprintf(stderr, "# DIRTY retries exhausted %s\n", cell_label);
+            res.trusted = false;
+            return res;
+        }
+
+        const size_t dirty_from = window_start();
+
+        const double a_cool =
+            cool_until_steady(backend, cell, reps, set_cand, clear_cand, baseline_cand, anchor_ref, cool, cell_label);
+        if (a_cool <= 0.0) {
+            res.trusted = false;
+            return res;
+        }
+
+        // the converged anchor is the only clean one now; re-measure the dirty window from it
+        anchors.clear();
+        anchors.push_back({ a_cool, dirty_from });
+
+        retries_left--;
+
+        fprintf(stderr, "# REDO candidates %zu..%zu %s\n", dirty_from, i, cell_label);
+        for (size_t j = dirty_from; j <= i; ++j) {
+            res.t[order[j]] = measure_one(backend, cell, reps, set_cand, clear_cand, order[j]);
+        }
+    }
+
+    return res;
+}

tools/tuning/bench.h

@@ -0,0 +1,57 @@
+#pragma once
+
+#include "ggml-backend.h"
+#include "ggml-cpp.h"
+#include "ggml.h"
+
+#include <cstdint>
+#include <functional>
+#include <vector>
+
+// A prebuilt graph replicated to amortize dispatch and synchronization overhead.
+struct perf_cell {
+    ggml_context_ptr        ctx;
+    ggml_backend_buffer_ptr buf;
+    ggml_cgraph *           gf     = nullptr;
+    int                     n_runs = 0;
+};
+
+using build_graph_fn  = std::function<ggml_tensor *(ggml_context *)>;
+using init_tensors_fn = std::function<void(ggml_context *)>;
+using op_flops_fn     = std::function<uint64_t(ggml_tensor *)>;
+
+perf_cell build_perf_cell(ggml_backend_t          backend,
+                          const build_graph_fn &  build,
+                          const init_tensors_fn & init,
+                          const op_flops_fn &     flops);
+
+double time_cell_median(ggml_backend_t backend, const perf_cell & cell, int reps);
+
+struct cooldown_opts {
+    bool   enabled   = true;
+    double drift     = 0.10;  // anchor drift that triggers a cooldown
+    double eps       = 0.03;  // anchor tolerance to call the GPU cool again
+    int    max_wait  = 120;  // seconds of cooling per cell before giving up
+    int    max_retry = 2;    // re-measure rounds per cell before giving up
+};
+
+using set_candidate_fn   = std::function<void(int)>;
+using clear_candidate_fn = std::function<void()>;
+
+struct cell_result {
+    std::vector<double> t;
+    bool                trusted    = true;
+    double              anchor_min = 0.0;
+    double              anchor_max = 0.0;
+};
+
+// Times candidates in order while using baseline_cand as a thermal-drift anchor.
+cell_result measure_cell(ggml_backend_t             backend,
+                         const perf_cell &          cell,
+                         int                        reps,
+                         const std::vector<int> &   order,
+                         const set_candidate_fn &   set_cand,
+                         const clear_candidate_fn & clear_cand,
+                         int                        baseline_cand,
+                         const cooldown_opts &      cool,
+                         const char *               cell_label);

tools/tuning/fa-vec.cpp

@@ -0,0 +1,641 @@
+#include "fa-vec.h"
+
+#include "bench.h"
+#include "ggml-backend.h"
+#include "ggml-metal-tuning.h"
+#include "ggml.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstdio>
+#include <cstring>
+#include <random>
+#include <set>
+#include <string>
+#include <vector>
+
+// GQA spec-decode shape: enough query heads to keep the GPU busy so the Q>1 K/V-reuse
+// benefit is visible. nh KV heads, nr2 query heads each, nr3 batches.
+static const int FA_NH  = 4;
+static const int FA_NR2 = 8;
+static const int FA_NR3 = 1;
+
+struct fa_shape {
+    int       dk;
+    int       dv;
+    int       ne01;  // query rows
+    int       ne11;  // KV length
+    ggml_type type_kv;
+};
+
+// mirrors test_flash_attn_ext::build_graph for the subset this tuner sweeps
+// (mask=true, sinks=false, prec=F32, type_K==type_V, no permute)
+static ggml_tensor * fa_build_graph(ggml_context * ctx, const fa_shape & s) {
+    const int64_t dk_padded = GGML_PAD(s.dk, ggml_blck_size(s.type_kv));
+    const int64_t dv_padded = GGML_PAD(s.dv, ggml_blck_size(s.type_kv));
+
+    ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, dk_padded, s.ne01, FA_NH * FA_NR2, FA_NR3);
+    ggml_set_name(q, "q");
+
+    // K/V are views of a 2x-tall parent, as they are of the KV cache in production
+    ggml_tensor * k0 = ggml_new_tensor_4d(ctx, s.type_kv, dk_padded, 2 * s.ne11, FA_NH, FA_NR3);
+    ggml_tensor * k  = ggml_view_4d(ctx, k0, dk_padded, s.ne11, FA_NH, FA_NR3, k0->nb[1], k0->nb[2], k0->nb[3], 0);
+    ggml_set_name(k, "k");
+
+    ggml_tensor * v = nullptr;
+    if (dk_padded == 576 && dv_padded == 512) {
+        // MLA: the V cache is a sub-view of the K cache
+        v = ggml_view_4d(ctx, k, dv_padded, s.ne11, FA_NH, FA_NR3, k->nb[1], k->nb[2], k->nb[3], 0);
+    } else {
+        ggml_tensor * v0 = ggml_new_tensor_4d(ctx, s.type_kv, dv_padded, 2 * s.ne11, FA_NH, FA_NR3);
+        v                = ggml_view_4d(ctx, v0, dv_padded, s.ne11, FA_NH, FA_NR3, v0->nb[1], v0->nb[2], v0->nb[3], 0);
+    }
+    ggml_set_name(v, "v");
+
+    ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, s.ne11, s.ne01, 1, FA_NR3);
+    ggml_set_name(m, "m");
+
+    ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f / sqrtf((float) s.dk), 0.0f, 0.0f);
+    ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32);
+    ggml_set_name(out, "out");
+
+    return out;
+}
+
+static uint64_t fa_op_flops(const fa_shape & s) {
+    // Q*K^T is ne01 x dk x ne11, P*V is ne01 x ne11 x dv, per head
+    return (uint64_t) 2 * FA_NH * FA_NR2 * s.ne01 * (s.dk + s.dv) * s.ne11 * FA_NR3;
+}
+
+static void fa_init_uniform(ggml_tensor * t, std::mt19937 & rng, float min, float max) {
+    const size_t nels = ggml_nelements(t);
+
+    std::vector<float>                    data(nels);
+    std::uniform_real_distribution<float> dist(min, max);
+    for (size_t i = 0; i < nels; i++) {
+        data[i] = dist(rng);
+    }
+
+    if (t->type == GGML_TYPE_F32) {
+        ggml_backend_tensor_set(t, data.data(), 0, nels * sizeof(float));
+        return;
+    }
+
+    GGML_ASSERT(ggml_is_quantized(t->type) || t->type == GGML_TYPE_F16 || t->type == GGML_TYPE_BF16);
+    GGML_ASSERT(nels % ggml_blck_size(t->type) == 0);
+
+    std::vector<float> imatrix(t->ne[0], 1.0f);
+    const float *      im = imatrix.data();
+    if (!ggml_quantize_requires_imatrix(t->type)) {
+        // when the imatrix is optional, exercise both paths; pick via one of the random numbers
+        if (data[0] > 0.5f * (min + max)) {
+            im = nullptr;
+        }
+    }
+
+    const size_t blck_size = ggml_blck_size(t->type);
+    const size_t n_blocks  = nels / blck_size;
+
+    std::vector<uint8_t> dataq(ggml_row_size(t->type, nels));
+    ggml_quantize_chunk(t->type, data.data(), dataq.data(), 0, n_blocks, blck_size, im);
+
+    ggml_backend_tensor_set(t, dataq.data(), 0, dataq.size());
+}
+
+// mirrors init_tensor_kq_mask: f16 mask with ~20% of its blocks set to -INF or zero.
+// the -INF blocks are what drives the kernel's skip-INF path, so this pattern is
+// load-bearing for the timings, not just for numerics.
+static void fa_init_kq_mask(ggml_tensor * t, std::mt19937 & rng, float min, float max) {
+    GGML_ASSERT(t->type == GGML_TYPE_F16);
+
+    const int32_t ne0 = (int32_t) t->ne[0];
+    const int32_t ne1 = (int32_t) t->ne[1];
+    const int32_t ne2 = (int32_t) t->ne[2];
+    const int32_t ne3 = (int32_t) t->ne[3];
+
+    std::vector<float>       data_f32(size_t(ne0) * ne1 * ne2 * ne3);
+    std::vector<ggml_fp16_t> data_f16(size_t(ne0) * ne1 * ne2 * ne3);
+
+    std::uniform_real_distribution<float> dis(min, max);
+    for (size_t i = 0; i < data_f32.size(); i++) {
+        data_f32[i] = dis(rng);
+    }
+
+    const int blck0 = 128;
+    const int blck1 = 64;
+
+    const int n_inf_zero_blocks = 0.2 * (ne0 * ne1 * ne2 * ne3) / (blck0 * blck1);
+
+    for (int b = 0; b < n_inf_zero_blocks; b++) {
+        const int p3 = (int) (rng() % ne3);
+        const int p2 = (int) (rng() % ne2);
+        const int p1 = (int) (rng() % ne1);
+        const int p0 = (int) (rng() % ne0);
+
+        const bool inf = rng() & 1;
+
+        for (int i1 = 0; i1 < blck1 && p1 + i1 < ne1; i1++) {
+            const int idx = p3 * ne2 * ne1 * ne0 + p2 * ne1 * ne0 + (p1 + i1) * ne0 + p0;
+
+            for (int i0 = 0; i0 < blck0 && p0 + i0 < ne0; i0++) {
+                data_f32[idx + i0] = inf ? -INFINITY : 0.0f;
+            }
+        }
+    }
+
+    ggml_fp32_to_fp16_row(data_f32.data(), data_f16.data(), ne0 * ne1 * ne2 * ne3);
+
+    ggml_backend_tensor_set(t, data_f16.data(), 0, data_f16.size() * sizeof(ggml_fp16_t));
+}
+
+static unsigned fa_cell_seed(const fa_shape & s, unsigned base) {
+    unsigned h = base;
+    for (int v : { s.dk, s.dv, s.ne01, s.ne11, (int) s.type_kv }) {
+        h = h * 1000003u + (unsigned) v;
+    }
+    return h;
+}
+
+static void fa_init_tensors(ggml_context * ctx, const fa_shape & s, unsigned base_seed) {
+    std::mt19937 rng(fa_cell_seed(s, base_seed));
+
+    for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+        if (t->view_src != NULL) {
+            continue;  // views share their parent's data
+        }
+        if (strcmp(t->name, "m") == 0) {
+            fa_init_kq_mask(t, rng, -1.0f, 1.0f);
+        } else {
+            fa_init_uniform(t, rng, -1.0f, 1.0f);
+        }
+    }
+}
+
+using set_override_t   = void (*)(int, int);
+using clear_override_t = void (*)(void);
+using bucket_t         = int (*)(int64_t);
+using baseline_ne_t    = int (*)(int, int);
+using device_token_t   = const char * (*) (ggml_backend_dev_t);
+
+struct fa_procs {
+    set_override_t   set_ov      = nullptr;
+    clear_override_t clr_ov      = nullptr;
+    bucket_t         ne11_bucket = nullptr;
+    bucket_t         ne01_bucket = nullptr;
+    baseline_ne_t    baseline_ne = nullptr;
+    device_token_t   dev_token   = nullptr;
+
+    bool ok() const { return set_ov && clr_ov && ne11_bucket && ne01_bucket && baseline_ne && dev_token; }
+};
+
+static fa_procs fa_resolve_procs(ggml_backend_dev_t dev) {
+    ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
+
+    fa_procs p;
+    p.set_ov = (set_override_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_metal_tuning_set_fa_vec_override");
+    p.clr_ov =
+        (clear_override_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_metal_tuning_clear_fa_vec_override");
+    p.ne11_bucket = (bucket_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_metal_tuning_fa_vec_ne11_bucket");
+    p.ne01_bucket = (bucket_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_metal_tuning_fa_vec_ne01_bucket");
+    p.baseline_ne =
+        (baseline_ne_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_metal_tuning_fa_vec_baseline_ne");
+    p.dev_token = (device_token_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_metal_tuning_device_token");
+
+    return p;
+}
+
+static bool fa_filter_has(const char * filter, const char * name) {
+    if (!filter) {
+        return true;
+    }
+
+    const std::string f = std::string(",") + filter + ",";
+
+    return f.find(std::string(",") + name + ",") != std::string::npos;
+}
+
+struct fa_cand {
+    int Q, NE;
+};
+
+struct fa_point {
+    int                 dk, dv, ne11, ne01;
+    std::vector<double> t;
+};
+
+// base_i identifies the (Q=1, baseline NE) anchor configuration.
+static std::vector<fa_cand> fa_build_cands(const fa_procs & procs, int dk, int dv, int & base_i) {
+    const int base_ne = procs.baseline_ne(dk, dv);
+
+    std::vector<fa_cand> cands;
+    base_i = -1;
+    for (int ne : ggml_metal_tuning::fa_vec_legal_ne(dk, dv)) {
+        for (int Q : { 1, 2, 4 }) {
+            if (Q == 1 && ne == base_ne) {
+                base_i = (int) cands.size();
+            }
+            cands.push_back({ Q, ne });
+        }
+    }
+    GGML_ASSERT(base_i >= 0);
+
+    return cands;
+}
+
+bool tuner_fa_vec_run(ggml_backend_t backend, ggml_backend_dev_t dev, const tuner_opts & opts) {
+    const fa_procs procs = fa_resolve_procs(dev);
+    if (!procs.ok()) {
+        fprintf(stderr, "error: metal fa_vec tuning procs unavailable\n");
+        return false;
+    }
+
+    const char * dev_token = procs.dev_token(dev);
+
+    struct shape_t {
+        int dk, dv;
+    };
+
+    const shape_t shapes[] = {
+        { 32,  32  },
+        { 64,  64  },
+        { 96,  96  },
+        { 128, 128 },
+        { 192, 192 },
+        { 192, 128 },
+        { 256, 256 },
+        { 320, 256 },
+        { 512, 512 },
+        { 576, 512 }
+    };
+    // nsg is a pipeline specialization constant (1 up to ne11=2048, 2 up to 4096, 4 above), so ne11
+    // bucket 1 takes two samples to cover both of its regimes. Bucket 0 is not sampled at all: the
+    // runtime leaves short KV at baseline, so no measurement there can reach the table.
+    const int ne11_rep[] = { 2048, 3072, 8192, 32768 };
+    const int ne01_rep[] = { 1, 2, 3, 4, 5, 6, 7, 8, 16 };  // point buckets (1-4) + tail mod-4 cycle + anchor
+
+    struct dtype_t {
+        ggml_type    type;
+        const char * token;
+    };
+
+    const dtype_t dtypes[] = {
+        { GGML_TYPE_F16,  "GGML_TYPE_F16"  },
+        { GGML_TYPE_Q4_0, "GGML_TYPE_Q4_0" },
+        { GGML_TYPE_Q4_1, "GGML_TYPE_Q4_1" },
+        { GGML_TYPE_Q5_0, "GGML_TYPE_Q5_0" },
+        { GGML_TYPE_Q5_1, "GGML_TYPE_Q5_1" },
+        { GGML_TYPE_Q8_0, "GGML_TYPE_Q8_0" },
+    };
+
+    const double TUNE_TAU   = 0.05;  // max POINTWISE regret to ride a domain default
+    const double TUNE_THETA = 1.05;  // min AGGREGATE bucket speedup vs baseline to tune at all
+
+    const cooldown_opts cool = {
+        opts.cooldown, opts.cool_drift, opts.cool_eps, opts.cool_max_wait, opts.cool_max_retry,
+    };
+
+    fprintf(stderr, "seed=%u reps=%d cooldown=%s (drift=%.2f eps=%.2f max_wait=%ds max_retry=%d)\n", opts.seed,
+            opts.reps, cool.enabled ? "on" : "off", cool.drift, cool.eps, cool.max_wait, cool.max_retry);
+    fprintf(stderr, "device token: %s\n", dev_token);
+
+    int n_untrusted = 0;
+
+    // stdout carries nothing but table rows, so the whole stream pastes into fa_vec_tuned_table
+    for (const auto & dtype : dtypes) {
+        const ggml_type type_kv = dtype.type;
+        if (!fa_filter_has(opts.dtype_filter, ggml_type_name(type_kv))) {
+            continue;
+        }
+
+        fprintf(stderr, "\n### dtype=%s\n", ggml_type_name(type_kv));
+
+        std::vector<fa_point> pts;
+
+        for (auto s : shapes) {
+            if (!fa_filter_has(opts.dk_filter, std::to_string(s.dk).c_str())) {
+                continue;
+            }
+
+            int                  base_i = 0;
+            std::vector<fa_cand> cands  = fa_build_cands(procs, s.dk, s.dv, base_i);
+
+            for (int ne11 : ne11_rep) {
+                for (int ne01 : ne01_rep) {
+                    const fa_shape sh = { s.dk, s.dv, ne01, ne11, type_kv };
+
+                    perf_cell cell = build_perf_cell(
+                        backend, [&](ggml_context * ctx) { return fa_build_graph(ctx, sh); },
+                        [&](ggml_context * ctx) { fa_init_tensors(ctx, sh, opts.seed); },
+                        [&](ggml_tensor *) { return fa_op_flops(sh); });
+
+                    if (cell.gf == nullptr) {
+                        continue;
+                    }
+
+                    // randomize candidate order to decorrelate thermal drift across the cell
+                    std::vector<int> order((size_t) cands.size());
+                    for (size_t i = 0; i < order.size(); ++i) {
+                        order[i] = (int) i;
+                    }
+                    std::shuffle(order.begin(), order.end(), std::mt19937(fa_cell_seed(sh, opts.seed)));
+
+                    char label[128];
+                    snprintf(label, sizeof(label), "dk=%d ne11=%d", s.dk, ne11);
+
+                    cell_result r = measure_cell(
+                        backend, cell, opts.reps, order, [&](int i) { procs.set_ov(cands[i].Q, cands[i].NE); },
+                        [&]() { procs.clr_ov(); }, base_i, cool, label);
+
+                    if (r.anchor_min > 0.0) {
+                        fprintf(stderr, "# noise dk=%d dv=%d ne11=%d ne01=%d spread=%.1f%%\n", s.dk, s.dv, ne11, ne01,
+                                100.0 * (r.anchor_max - r.anchor_min) / r.anchor_min);
+                    }
+
+                    if (!r.trusted) {
+                        n_untrusted++;
+                        fprintf(stderr, "# DROP untrusted cell dk=%d dv=%d ne11=%d ne01=%d\n", s.dk, s.dv, ne11, ne01);
+                        continue;
+                    }
+
+                    int best_i = -1;
+                    for (size_t i = 0; i < cands.size(); ++i) {
+                        if (r.t[i] > 0.0 && (best_i < 0 || r.t[i] < r.t[best_i])) {
+                            best_i = (int) i;
+                        }
+                    }
+                    const double base_t = r.t[base_i];
+                    const bool   keep   = best_i >= 0 && base_t > 0.0 && r.t[best_i] < base_t * 0.98;
+
+                    fprintf(stderr, "# dtype=%s dk=%d dv=%d ne11=%d ne01=%d:", ggml_type_name(type_kv), s.dk, s.dv,
+                            ne11, ne01);
+                    for (size_t i = 0; i < cands.size(); ++i) {
+                        fprintf(stderr, "  Q%dNE%d=%.1f%s", cands[i].Q, cands[i].NE, r.t[i],
+                                (int) i == best_i ? "*" : "");
+                    }
+                    if (keep) {
+                        fprintf(stderr, "  => Q%d,NE%d  %.2fx\n", cands[best_i].Q, cands[best_i].NE,
+                                base_t / r.t[best_i]);
+                    } else {
+                        fprintf(stderr, "  => baseline\n");
+                    }
+
+                    pts.push_back({ s.dk, s.dv, ne11, ne01, r.t });
+                }
+            }
+        }
+
+        // compress into pasteable rows. per (dk,dv) and ne01 domain {decode==1, batch>=2},
+        // emit one ne11-collapsed default cfg (ne11_b=-1) plus a per-bucket exception wherever the
+        // default's pointwise regret vs the bucket target exceeds TUNE_TAU, or the default is not
+        // admissible for that bucket (see never_slower / admissible below).
+        std::vector<std::string> rows_out;
+        char                     rbuf[192];
+
+        for (auto s : shapes) {
+            if (!fa_filter_has(opts.dk_filter, std::to_string(s.dk).c_str())) {
+                continue;
+            }
+
+            int                  base_i = 0;
+            std::vector<fa_cand> cands  = fa_build_cands(procs, s.dk, s.dv, base_i);
+
+            struct bkt_t {
+                int                           b11, b01, Ti;
+                std::vector<double>           agg;
+                std::vector<const fa_point *> bp;
+            };
+
+            // A config may represent a bucket only if it is no slower than baseline at every point that
+            // bucket covers. The aggregate gate below sums absolute times, so it can pass on the aligned
+            // and deep points while a misaligned ne01 pays the mod-Q padding. Nothing measured, nothing
+            // proven: a bucket with no surviving sample admits baseline only.
+            auto never_slower = [&](const std::vector<const fa_point *> & bp, int i) {
+                if (i == base_i) {
+                    return true;
+                }
+                if (bp.empty()) {
+                    return false;
+                }
+                for (const auto * p : bp) {
+                    if (p->t[i] <= 0.0 || p->t[base_i] <= 0.0 || p->t[i] > p->t[base_i]) {
+                        return false;
+                    }
+                }
+                return true;
+            };
+
+            // The padded-row waste ceil(n/Q)*Q/n is largest at the smallest ne01 of each residue class
+            // mod Q, so one of a bucket's first Q values carries the worst padding it can ever see, and
+            // that value has to be sampled. Otherwise the bucket bounds nothing: a config picked on the
+            // aligned ne01=8,16 says nothing about ne01=9. This covers the padding term only - the
+            // per-row cost varies with ne01 too - so it is a floor on the evidence, not a proof.
+            auto admissible = [&](const std::vector<const fa_point *> & bp, int b01, int i) {
+                if (!never_slower(bp, i)) {
+                    return false;
+                }
+                const int Q = cands[i].Q;
+                if (Q == 1) {
+                    return true;  // one row per threadgroup, no padding to witness
+                }
+                int lo = bp[0]->ne01;
+                for (const auto * p : bp) {
+                    lo = std::min(lo, p->ne01);
+                }
+                while (lo > 1 && procs.ne01_bucket(lo - 1) == b01) {
+                    lo--;  // walk down to where this bucket's runtime domain starts
+                }
+                int    wit  = lo;
+                double wmax = 0.0;
+                for (int n = lo; n < lo + Q && procs.ne01_bucket(n) == b01; ++n) {
+                    const int    padded = ((n + Q - 1) / Q) * Q;
+                    const double w      = (double) padded / n;
+                    if (w > wmax) {
+                        wmax = w;
+                        wit  = n;
+                    }
+                }
+                for (const auto * p : bp) {
+                    if (p->ne01 == wit) {
+                        return true;
+                    }
+                }
+                return false;
+            };
+
+            std::set<std::pair<int, int>> buckets;
+            for (int ne11 : ne11_rep) {
+                const int b11 = procs.ne11_bucket(ne11);
+                if (b11 == 0) {
+                    continue;
+                }
+                for (int ne01 : ne01_rep) {
+                    buckets.insert({ b11, procs.ne01_bucket(ne01) });
+                }
+            }
+
+            std::vector<bkt_t> bks;
+            for (const auto & bb : buckets) {
+                const int b11 = bb.first, b01 = bb.second;
+
+                std::vector<const fa_point *> bp;
+                for (const auto & p : pts) {
+                    if (p.dk == s.dk && p.dv == s.dv && procs.ne11_bucket(p.ne11) == b11 &&
+                        procs.ne01_bucket(p.ne01) == b01) {
+                        bp.push_back(&p);
+                    }
+                }
+
+                fprintf(stderr, "# bucket dk=%d dv=%d ne11_b=%d ne01_b=%d samples=%zu\n", s.dk, s.dv, b11, b01,
+                        bp.size());
+                if (bp.empty()) {
+                    // nothing to check a config against, so pin the bucket to baseline instead of
+                    // letting the ne11-collapsed domain default ride in unmeasured
+                    fprintf(stderr, "# WARN empty bucket dk=%d dv=%d ne11_b=%d ne01_b=%d -> baseline\n", s.dk, s.dv,
+                            b11, b01);
+                    bks.push_back({ b11, b01, base_i, std::vector<double>(cands.size(), 0.0), {} });
+                    continue;
+                }
+
+                std::vector<double> agg(cands.size(), 0.0), worst(cands.size(), 0.0);
+                for (const auto * p : bp) {
+                    double bestt = 0.0;
+                    for (size_t i = 0; i < cands.size(); ++i) {
+                        if (p->t[i] > 0.0 && (bestt == 0.0 || p->t[i] < bestt)) {
+                            bestt = p->t[i];
+                        }
+                    }
+                    for (size_t i = 0; i < cands.size(); ++i) {
+                        agg[i] += p->t[i];
+                        if (p->t[i] > 0.0 && bestt > 0.0) {
+                            worst[i] = std::max(worst[i], p->t[i] / bestt);
+                        }
+                    }
+                }
+
+                int robust = -1, oracle_pick = -1;
+                for (size_t i = 0; i < cands.size(); ++i) {
+                    auto tighter = [&](int j) {
+                        return j < 0 || worst[i] < worst[j] ||
+                               (worst[i] == worst[j] && (cands[i].Q < cands[j].Q ||
+                                                         (cands[i].Q == cands[j].Q && cands[i].NE < cands[j].NE)));
+                    };
+                    if (tighter(oracle_pick)) {
+                        oracle_pick = (int) i;
+                    }
+                    if (admissible(bp, b01, (int) i) && tighter(robust)) {
+                        robust = (int) i;
+                    }
+                }
+
+                const bool tune = robust != base_i && agg[base_i] > 0.0 && agg[robust] > 0.0 &&
+                                  agg[base_i] / agg[robust] >= TUNE_THETA;
+
+                // report what the no-harm rule cost this bucket, but only when it changed the outcome:
+                // a sweep on another machine then shows where the winner loses, instead of just
+                // emitting a smaller table
+                const bool refused = oracle_pick != robust && oracle_pick != base_i && agg[base_i] > 0.0 &&
+                                     agg[oracle_pick] > 0.0 && agg[base_i] / agg[oracle_pick] >= TUNE_THETA;
+                if (refused) {
+                    double over = 0.0;
+                    int    at11 = 0, at01 = 0;
+                    for (const auto * p : bp) {
+                        if (p->t[base_i] > 0.0 && p->t[oracle_pick] / p->t[base_i] - 1.0 > over) {
+                            over = p->t[oracle_pick] / p->t[base_i] - 1.0;
+                            at11 = p->ne11;
+                            at01 = p->ne01;
+                        }
+                    }
+                    if (over > 0.0) {
+                        fprintf(stderr,
+                                "# reject dk=%d dv=%d ne11_b=%d ne01_b=%d Q%dNE%d: +%.2f%% vs baseline at "
+                                "ne11=%d ne01=%d\n",
+                                s.dk, s.dv, b11, b01, cands[oracle_pick].Q, cands[oracle_pick].NE, 100.0 * over, at11,
+                                at01);
+                    } else {
+                        fprintf(stderr, "# reject dk=%d dv=%d ne11_b=%d ne01_b=%d Q%dNE%d: no padding witness\n", s.dk,
+                                s.dv, b11, b01, cands[oracle_pick].Q, cands[oracle_pick].NE);
+                    }
+                }
+
+                bks.push_back({ b11, b01, tune ? robust : base_i, agg, bp });
+            }
+
+            // pointwise regret of default cfg d vs the bucket target: a ratio-of-sums lets a
+            // default that wins on aligned ne01 hide a large penalty on a misaligned point
+            auto reg_pointwise = [&](const bkt_t * b, int d) {
+                double r = 0.0;
+                for (const auto * p : b->bp) {
+                    const double td = p->t[d], tT = p->t[b->Ti];
+                    if (td > 0.0 && tT > 0.0) {
+                        r = std::max(r, td / tT - 1.0);
+                    }
+                }
+                return r;
+            };
+
+            for (int dom = 0; dom <= 1; ++dom) {  // 0 = decode (ne01==1), 1 = batch (ne01>=2)
+                std::vector<const bkt_t *> db;
+                for (const auto & b : bks) {
+                    if ((dom == 0) == (b.b01 == 0)) {
+                        db.push_back(&b);
+                    }
+                }
+                if (db.empty()) {
+                    continue;
+                }
+
+                // default cfg = the one minimizing (#rows, total achieved time, Q, NE)
+                int    bestD = -1, bestRows = 1 << 30;
+                double bestTot = 0.0;
+                for (size_t d = 0; d < cands.size(); ++d) {
+                    int    rows = ((int) d != base_i) ? 1 : 0;
+                    double tot  = 0.0;
+                    for (const auto * b : db) {
+                        if (reg_pointwise(b, (int) d) > TUNE_TAU || !admissible(b->bp, b->b01, (int) d)) {
+                            rows++;
+                            tot += b->agg[b->Ti];
+                        } else {
+                            tot += b->agg[d];
+                        }
+                    }
+                    const bool better =
+                        bestD < 0 || rows < bestRows ||
+                        (rows == bestRows &&
+                         (tot < bestTot ||
+                          (tot == bestTot && (cands[d].Q < cands[bestD].Q ||
+                                              (cands[d].Q == cands[bestD].Q && cands[d].NE < cands[bestD].NE)))));
+                    if (better) {
+                        bestD    = (int) d;
+                        bestRows = rows;
+                        bestTot  = tot;
+                    }
+                }
+
+                if (bestD != base_i) {
+                    snprintf(rbuf, sizeof(rbuf), "    { { %s, %s, %d, %d, -1, %d }, { %d, %d } },", dev_token,
+                             dtype.token, s.dk, s.dv, dom, cands[bestD].Q, cands[bestD].NE);
+                    rows_out.emplace_back(rbuf);
+                }
+                for (const auto * b : db) {
+                    if (reg_pointwise(b, bestD) <= TUNE_TAU && admissible(b->bp, b->b01, bestD)) {
+                        continue;
+                    }
+                    snprintf(rbuf, sizeof(rbuf), "    { { %s, %s, %d, %d, %d, %d }, { %d, %d } },", dev_token,
+                             dtype.token, s.dk, s.dv, b->b11, b->b01, cands[b->Ti].Q, cands[b->Ti].NE);
+                    rows_out.emplace_back(rbuf);
+                }
+            }
+        }
+
+        for (const auto & r : rows_out) {
+            printf("%s\n", r.c_str());
+        }
+        fflush(stdout);
+    }
+
+    if (n_untrusted > 0) {
+        fprintf(stderr, "\n%d cells excluded as untrusted (see DROP lines above)\n", n_untrusted);
+    }
+
+    return true;
+}

tools/tuning/fa-vec.h

@@ -0,0 +1,18 @@
+#pragma once
+
+#include "ggml-backend.h"
+
+struct tuner_opts {
+    const char * dtype_filter   = nullptr;  // comma-separated, e.g. "f16,q4_0"; null = all
+    const char * dk_filter      = nullptr;  // comma-separated dk values, e.g. "128,192"; null = all
+    int          reps           = 7;
+    unsigned     seed           = 1234;
+    bool         cooldown       = true;
+    double       cool_drift     = 0.10;
+    double       cool_eps       = 0.03;
+    int          cool_max_wait  = 120;
+    int          cool_max_retry = 2;
+};
+
+// Returns false only when the required Metal proc bridges are unavailable.
+bool tuner_fa_vec_run(ggml_backend_t backend, ggml_backend_dev_t dev, const tuner_opts & opts);

tools/tuning/main.cpp

@@ -0,0 +1,139 @@
+#include "fa-vec.h"
+#include "ggml-backend.h"
+#include "ggml.h"
+
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+
+struct tuner_def {
+    const char * name;
+    bool (*run)(ggml_backend_t, ggml_backend_dev_t, const tuner_opts &);
+};
+
+static const tuner_def k_tuners[] = {
+    { "fa-vec", tuner_fa_vec_run },
+};
+
+static void usage(const char * argv0) {
+    printf("usage: %s <tuner> [options]\n", argv0);
+    printf("\n");
+    printf("  offline kernel tuner for the Metal backend: sweeps a kernel's config grid and\n");
+    printf("  prints pasteable table rows for the machine it runs on. never a pass/fail test.\n");
+    printf("\n");
+    printf("  tuners:\n");
+    printf("    fa-vec            flash-attn vec (Q,NE) for ggml-metal-tuning.cpp\n");
+    printf("\n");
+    printf("  options:\n");
+    printf("    -b <name>         backend device (default: first Metal device)\n");
+    printf("    --dtype <list>    restrict KV dtypes, e.g. f16,q4_0 (default: all)\n");
+    printf("    --dk <list>       restrict head sizes, e.g. 128,192 (default: all)\n");
+    printf("    --reps <n>        timed reps per candidate, odd for an exact median (default: 7)\n");
+    printf("    --seed <n>        RNG seed; per-cell seeds mix it with the shape (default: 1234)\n");
+    printf("    --no-cooldown     do not pause/re-measure on thermal drift, only warn\n");
+    printf("    --cool-drift <f>  anchor drift that triggers a cooldown (default: 0.10)\n");
+    printf("    --cool-eps <f>    anchor tolerance to consider the GPU cool again (default: 0.03)\n");
+    printf("    --cool-max-wait <s> give up cooling a cell after this many seconds (default: 120)\n");
+    printf("    --cool-max-retry <n> re-measure rounds per cell before giving up (default: 2)\n");
+    printf("\n");
+    printf("  the table goes to stdout, all diagnostics to stderr:\n");
+    printf("    %s fa-vec > rows.txt 2> sweep.log\n", argv0);
+}
+
+int main(int argc, char ** argv) {
+    const char * tuner = nullptr;
+    const char * bname = nullptr;
+    tuner_opts   opts;
+
+    for (int i = 1; i < argc; i++) {
+        const char * a = argv[i];
+        if (strcmp(a, "-h") == 0 || strcmp(a, "--help") == 0) {
+            usage(argv[0]);
+            return 0;
+        } else if (strcmp(a, "-b") == 0 && i + 1 < argc) {
+            bname = argv[++i];
+        } else if (strcmp(a, "--dtype") == 0 && i + 1 < argc) {
+            opts.dtype_filter = argv[++i];
+        } else if (strcmp(a, "--dk") == 0 && i + 1 < argc) {
+            opts.dk_filter = argv[++i];
+        } else if (strcmp(a, "--reps") == 0 && i + 1 < argc) {
+            opts.reps = atoi(argv[++i]);
+        } else if (strcmp(a, "--seed") == 0 && i + 1 < argc) {
+            opts.seed = (unsigned) strtoul(argv[++i], nullptr, 10);
+        } else if (strcmp(a, "--no-cooldown") == 0) {
+            opts.cooldown = false;
+        } else if (strcmp(a, "--cool-drift") == 0 && i + 1 < argc) {
+            opts.cool_drift = atof(argv[++i]);
+        } else if (strcmp(a, "--cool-eps") == 0 && i + 1 < argc) {
+            opts.cool_eps = atof(argv[++i]);
+        } else if (strcmp(a, "--cool-max-wait") == 0 && i + 1 < argc) {
+            opts.cool_max_wait = atoi(argv[++i]);
+        } else if (strcmp(a, "--cool-max-retry") == 0 && i + 1 < argc) {
+            opts.cool_max_retry = atoi(argv[++i]);
+        } else if (a[0] != '-' && tuner == nullptr) {
+            tuner = a;
+        } else {
+            fprintf(stderr, "error: unrecognized or incomplete argument: %s\n\n", a);
+            usage(argv[0]);
+            return 1;
+        }
+    }
+
+    if (tuner == nullptr) {
+        usage(argv[0]);
+        return 1;
+    }
+    if (opts.reps < 1) {
+        fprintf(stderr, "error: --reps must be >= 1\n");
+        return 1;
+    }
+
+    const tuner_def * t = nullptr;
+    for (const auto & cand : k_tuners) {
+        if (strcmp(tuner, cand.name) == 0) {
+            t = &cand;
+            break;
+        }
+    }
+    if (t == nullptr) {
+        fprintf(stderr, "error: unknown tuner: %s\n\n", tuner);
+        usage(argv[0]);
+        return 1;
+    }
+
+    ggml_backend_load_all();
+
+    ggml_backend_dev_t dev = nullptr;
+    for (size_t i = 0; i < ggml_backend_dev_count(); i++) {
+        ggml_backend_dev_t d = ggml_backend_dev_get(i);
+        if (bname) {
+            if (strcmp(ggml_backend_dev_name(d), bname) == 0) {
+                dev = d;
+                break;
+            }
+        } else if (strncmp(ggml_backend_dev_name(d), "MTL", 3) == 0) {
+            dev = d;
+            break;
+        }
+    }
+
+    if (dev == nullptr) {
+        fprintf(stderr, "error: no %s device found\n", bname ? bname : "Metal");
+        return 1;
+    }
+
+    ggml_backend_t backend = ggml_backend_dev_init(dev, nullptr);
+    if (backend == nullptr) {
+        fprintf(stderr, "error: failed to init backend %s\n", ggml_backend_dev_name(dev));
+        return 1;
+    }
+
+    fprintf(stderr, "device: %s (%s)\n", ggml_backend_dev_name(dev), ggml_backend_dev_description(dev));
+
+    const bool ok = t->run(backend, dev, opts);
+
+    ggml_backend_free(backend);
+    ggml_quantize_free();
+
+    return ok ? 0 : 1;
+}