metal : per-device tuned (Q, NE) for flash-attn vec (#25750)

* rebase Q-generic FA vec body from 01dc93607 (#23114)

* add 53 f16 (Q,NE) flash-attn vec instantiations (vec 80 -> 133)

* add FA vec (Q,NE) tuning table + dispatch wiring + SMEM cap fallback

* add  FA vec (Q,NE) perf sweep

* fill tuning result

* fold family table into a per-family representative SKU

* refactor tuning result format

* extend FA vec tuning to quantized KV caches

* sync fa vec tuner bucketing with runtime, use pointwise tuning regret

* update tuned table

* format and cleanup

* prefix fa_vec tuning procs with ggml_backend_metal_tuning_, drop unused fa_vec_override_active
This commit is contained in:
YiChen Lv
2026-07-29 21:46:01 +08:00
committed by GitHub
parent 7185fef97d
commit 96ece714c4
10 changed files with 1556 additions and 175 deletions

View File

@@ -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

View File

@@ -1,6 +1,7 @@
#include "ggml-metal-device.h"
#include "ggml-metal-impl.h"
#include "ggml-metal-tuning.h"
#include "ggml-impl.h"
@@ -1480,6 +1481,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) {
assert(op->op == GGML_OP_FLASH_ATTN_EXT);
@@ -1493,11 +1496,17 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
const int32_t ns10 = op->src[1]->nb[1]/op->src[1]->nb[0];
const int32_t ns20 = op->src[2]->nb[1]/op->src[2]->nb[0];
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",
ggml_type_name(op->src[1]->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,

View File

@@ -197,6 +197,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);
@@ -264,6 +266,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;
};

View File

@@ -1173,7 +1173,8 @@ ggml_metal_device_t ggml_metal_device_init(int device) {
{
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;
}
}

View File

@@ -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>
@@ -2942,12 +2943,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;
@@ -3006,7 +3013,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;
@@ -3026,6 +3033,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
}
ggml_metal_kargs_flash_attn_ext_vec args = {
/*.ne01 =*/ ne01,
/*.ne02 =*/ ne02,
@@ -3061,7 +3074,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);
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);
GGML_ASSERT(nsg*32 <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));

View File

@@ -0,0 +1,347 @@
#include "ggml-metal-tuning.h"
#include <cstddef>
#include <cstring>
#include <iterator>
namespace ggml_metal_tuning {
int fa_vec_ne11_bucket(int64_t ne11) {
for (int i = 0; i < (int) std::size(FA_VEC_NE11_BUCKETS); ++i) {
if (ne11 < FA_VEC_NE11_BUCKETS[i]) {
return i;
}
}
return (int) std::size(FA_VEC_NE11_BUCKETS);
}
int fa_vec_ne01_bucket(int64_t ne01) {
for (int i = 0; i < (int) std::size(FA_VEC_NE01_BUCKETS); ++i) {
if (ne01 < FA_VEC_NE01_BUCKETS[i]) {
return i;
}
}
return (int) std::size(FA_VEC_NE01_BUCKETS);
}
int fa_vec_baseline_ne(int dk, int dv) {
if (dk == 32 && dv == 32) {
return 4;
}
if (dk == 64 && dv == 64) {
return 2;
}
if (dk == 96 && dv == 96) {
return 4;
}
if (dk == 128 && dv == 128) {
return 1;
}
if (dk == 192 && dv == 192) {
return 2;
}
if (dk == 192 && dv == 128) {
return 2;
}
if (dk == 256 && dv == 256) {
return 1;
}
if (dk == 320 && dv == 256) {
return 2;
}
if (dk == 512 && dv == 512) {
return 1;
}
if (dk == 576 && dv == 512) {
return 2;
}
return 4; // template default
}
fa_vec_cfg_t fa_vec_baseline_cfg(int dk, int dv) {
return { 1, (int8_t) fa_vec_baseline_ne(dk, dv) };
}
// Generated by `test-backend-ops tune --tune-perf`; do not hand-edit.
// One row per kept bucket, plus per-(dtype,dk,dv) ne11-collapsed domain defaults
// (ne11_b = FA_VEC_NE11_DEFAULT, ne01_b = domain). To retune or add a device, re-run the
// sweep and paste its block. See ggml-metal-tuning.h for the row/lookup semantics.
constexpr fa_vec_entry_t fa_vec_tuned_table[] = {
// ---- f16: 13 rows ----
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 1, 4 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 2, 4 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 3, 3 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } },
// ---- q4_0: 29 rows ----
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 32, 32, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 32, 32, 3, 2 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 32, 32, 3, 3 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 32, 32, 3, 4 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 96, 96, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 96, 96, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 96, 96, 2, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 192, 128, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 192, 128, 3, 3 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 576, 512, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 576, 512, 1, 1 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 576, 512, 1, 2 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_0, 576, 512, 1, 4 }, { 1, 2 } },
// ---- q4_1: 28 rows ----
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 32, 32, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 32, 32, 3, 2 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 32, 32, 3, 3 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 32, 32, 3, 4 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 96, 96, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 96, 96, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 96, 96, 2, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 128, 128, 2, 3 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 128, 128, 2, 4 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 128, 128, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 128, 128, 3, 3 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 128, 128, 3, 4 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 192, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 192, 128, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 192, 128, 2, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 192, 128, 3, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 256, 256, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 256, 256, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 576, 512, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q4_1, 576, 512, -1, 1 }, { 1, 4 } },
// ---- q5_0: 45 rows ----
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 32, 32, 1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 32, 32, 2, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 32, 32, 2, 4 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 32, 32, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 96, 96, 1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 96, 96, 2, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 96, 96, 2, 4 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 96, 96, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 96, 96, 3, 4 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 192, 192, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 192, 192, 1, 1 }, { 2, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 192, 192, 1, 2 }, { 2, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 192, 192, 1, 3 }, { 2, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 2, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 256, 256, 1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 256, 256, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 256, 256, 1, 4 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 256, 256, 2, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 256, 256, 3, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 320, 256, -1, 1 }, { 2, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 320, 256, 1, 2 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 320, 256, 1, 3 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 320, 256, 1, 4 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 320, 256, 2, 2 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 320, 256, 3, 2 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 576, 512, 2, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 576, 512, 2, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 576, 512, 2, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 576, 512, 2, 3 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_0, 576, 512, 2, 4 }, { 1, 4 } },
// ---- q5_1: 49 rows ----
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 32, 32, 1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 32, 32, 2, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 32, 32, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 96, 96, 1, 2 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 96, 96, 1, 3 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 96, 96, 2, 2 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 96, 96, 2, 3 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 192, 192, 1, 1 }, { 2, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 192, 192, 1, 2 }, { 2, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 192, 192, 1, 3 }, { 2, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 192, 192, 1, 4 }, { 2, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 2, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 256, 256, 1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 256, 256, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 256, 256, 1, 4 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 256, 256, 2, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 256, 256, 3, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 320, 256, -1, 1 }, { 2, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 320, 256, 1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 320, 256, 1, 2 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 320, 256, 1, 3 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 320, 256, 1, 4 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 320, 256, 2, 2 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 320, 256, 3, 2 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 512, 512, 1, 1 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 512, 512, 1, 2 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 512, 512, 1, 3 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 512, 512, 1, 4 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 576, 512, 2, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 576, 512, 2, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 576, 512, 2, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 576, 512, 2, 3 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q5_1, 576, 512, 2, 4 }, { 1, 4 } },
// ---- q8_0: 29 rows ----
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 32, 32, 3, 2 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 32, 32, 3, 3 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 32, 32, 3, 4 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 128, 128, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 128, 128, 3, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 192, 128, 1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 192, 128, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 192, 128, 1, 3 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 192, 128, 2, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 192, 128, 3, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
};
static enum ggml_metal_device_id fa_vec_family_representative(int gpu_family) {
switch (gpu_family) {
case 9: return GGML_METAL_DEVICE_M4_MAX;
default: return GGML_METAL_DEVICE_GENERIC;
}
}
static bool g_override_set = false;
static fa_vec_cfg_t g_override_cfg = { 1, 4 };
void fa_vec_set_override(fa_vec_cfg_t cfg) {
g_override_cfg = cfg;
g_override_set = true;
}
void fa_vec_clear_override() {
g_override_set = false;
}
static const fa_vec_cfg_t * find_cfg(const fa_vec_entry_t * tbl, size_t n, const fa_vec_key_t & k) {
for (size_t i = 0; i < n; ++i) {
if (memcmp(&tbl[i].key, &k, sizeof(k)) == 0) {
return &tbl[i].cfg;
}
}
return nullptr;
}
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) {
if (g_override_set) {
return g_override_cfg;
}
const fa_vec_cfg_t baseline = fa_vec_baseline_cfg(dk, dv);
const int ne11_b = fa_vec_ne11_bucket(ne11);
if (ne11_b == 0) {
return baseline; // short KV: attention is a small slice of the step, left to baseline
}
const int ne01_b = fa_vec_ne01_bucket(ne01);
fa_vec_key_t k{};
k.dtype = (int8_t) dtype;
k.dk = (int16_t) dk;
k.dv = (int16_t) dv;
// exact bucket, then the ne01 domain default (ne11 collapsed); tried under each device tier
auto lookup = [&](enum ggml_metal_device_id dev) -> const fa_vec_cfg_t * {
k.device_id = (int8_t) dev;
k.ne11_b = (int8_t) ne11_b;
k.ne01_b = (int8_t) ne01_b;
if (auto * c = find_cfg(fa_vec_tuned_table, std::size(fa_vec_tuned_table), k)) {
return c;
}
k.ne11_b = FA_VEC_NE11_DEFAULT;
k.ne01_b = (ne01_b == 0) ? FA_VEC_DOMAIN_DECODE : FA_VEC_DOMAIN_BATCH;
return find_cfg(fa_vec_tuned_table, std::size(fa_vec_tuned_table), k);
};
if (auto * c = lookup(device_id)) {
return *c;
}
// family fallback: retry under the family's representative SKU; none -> baseline
if (gpu_family > 0) {
const enum ggml_metal_device_id rep = fa_vec_family_representative(gpu_family);
if (rep != GGML_METAL_DEVICE_GENERIC) {
if (auto * c = lookup(rep)) {
return *c;
}
}
}
return baseline;
}
} // namespace ggml_metal_tuning

View File

@@ -0,0 +1,63 @@
#pragma once
#include "ggml-metal-device.h" // enum ggml_metal_device_id
#include "ggml.h"
#include <cstdint>
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 (run_fa_vec_tune_check
// exercises every (Q,NE), so a missing instantiation surfaces there).
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;
};
// 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

View File

@@ -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>
@@ -868,10 +869,46 @@ 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 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;
}
return NULL;

File diff suppressed because it is too large Load Diff

View File

@@ -478,6 +478,7 @@ enum test_mode {
MODE_PERF,
MODE_GRAD,
MODE_SUPPORT,
MODE_TUNE,
};
// Output format support similar to llama-bench
@@ -10111,8 +10112,482 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_from_file(const c
return test_cases;
}
// ---- FA vec (Q,NE) tuning: numerical correctness check + drift guard ----
// metal proc_address bridges (resolved by string, not symbol linkage)
using set_fa_vec_override_t = void (*)(int, int);
using clear_fa_vec_override_t = void (*)(void);
using fa_vec_bucket_t = int (*)(int64_t); // runtime ne11/ne01 bucketers, shared via proc bridge
using fa_vec_baseline_ne_t = int (*)(int, int); // runtime (dk,dv) -> baseline NE
// legal NE for a (dk,dv): NL = 32/NE, require (dk/4)%NL==0 && (dv/4)%NL==0
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;
}
// Baseline-path smoke test: with no override, fa_vec_pick returns (1, baseline_ne) and
// dispatches the no-suffix baseline kernel. Checks baseline correctness across all 10 shapes
// (NE is numerically transparent; the (Q,NE) kernels are covered by run_fa_vec_tune_check).
static bool run_fa_vec_drift_guard(ggml_backend_t backend_metal, ggml_backend_t backend_cpu) {
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} };
bool ok = true;
for (auto s : shapes) {
// no override + short KV (ne11 < FA_VEC_NE11_BUCKETS[0]) -> fa_vec_pick returns baseline;
// compare metal vs CPU-ref
test_flash_attn_ext tc(s.dk, s.dv, /*nh=*/4, {1,1}, /*kv=*/512, /*nb=*/8,
/*mask=*/true, /*sinks=*/false, 0.0f, 0.0f, GGML_PREC_F32,
GGML_TYPE_F16, GGML_TYPE_F16);
auto st = tc.eval(backend_metal, backend_cpu, "FLASH_ATTN_EXT", nullptr);
if (st == test_status_t::FAIL) {
printf("DRIFT/baseline FAIL dk=%d dv=%d\n", s.dk, s.dv);
ok = false;
}
}
return ok;
}
// Numerical gate: for each (dk,dv) x legal (Q,NE), force the override and compare metal
// vs CPU-ref. The ne01/ne11/mask/sinks points exercise the rebased body's padded rows
// (ne01 not a multiple of Q), per-qq sinks, skip-INF, kvpad, and the nsg-dependent shmem
// offsets / parallel-reduce stride (ne11 -> nsg 1/2/4).
static bool run_fa_vec_tune_check(ggml_backend_t backend_metal, ggml_backend_t backend_cpu) {
auto * reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend_metal));
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) { printf("metal fa_vec override proc unavailable\n"); return false; }
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} };
const int ne01_pts[] = { 1, 2, 3, 8, 17 }; // 1, Q+/-1, 2Q-1, prime; vec upper bound < 20
const int ne11_pts[] = { 512, 4096, 8192 }; // drives adaptive nsg 1/2/4
const int ne11_kvpad[] = { 4097, 8193 }; // has_kvpad (non-32-multiple kv)
const bool mask_pts[] = { true, false };
const bool sinks_pts[] = { false, true };
bool ok = true;
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 (bool mask : mask_pts) {
for (bool sinks : sinks_pts) {
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, sinks, 0.0f, 0.0f, GGML_PREC_F32,
GGML_TYPE_F16, GGML_TYPE_F16);
auto st = tc.eval(backend_metal, backend_cpu, "FLASH_ATTN_EXT", nullptr);
clear_ov();
if (st == test_status_t::FAIL) {
printf("FAIL dk=%d dv=%d Q=%d ne=%d ne01=%d ne11=%d mask=%d sinks=%d\n",
s.dk, s.dv, Q, ne, ne01, ne11, (int) mask, (int) sinks);
ok = false; n_fail++;
}
n_run++;
}
}
}
}
}
}
}
// kvpad pass (non-32-multiple kv); dk128/256 only to keep the case count down
for (auto s : shapes) {
if (!((s.dk == 128 && s.dv == 128) || (s.dk == 256 && s.dv == 256))) {
continue;
}
for (int ne : fa_vec_legal_ne(s.dk, s.dv)) {
for (int Q : {1, 2, 4}) {
for (int ne11 : ne11_kvpad) {
set_ov(Q, ne);
test_flash_attn_ext tc(s.dk, s.dv, /*nh=*/4, {1,1}, /*kv=*/ne11, /*nb=*/8,
/*mask=*/true, /*sinks=*/false, 0.0f, 0.0f, GGML_PREC_F32,
GGML_TYPE_F16, GGML_TYPE_F16);
auto st = tc.eval(backend_metal, backend_cpu, "FLASH_ATTN_EXT", nullptr);
clear_ov();
if (st == test_status_t::FAIL) {
printf("FAIL(kvpad) dk=%d dv=%d Q=%d ne=%d ne11=%d\n", s.dk, s.dv, Q, ne, ne11);
ok = false; n_fail++;
}
n_run++;
}
}
}
}
// quantized K/V numerical check: the rebased body's dequant path is Q-generic
// (dequant once, reuse across Q rows); confirm it stays correct for every quant precision.
const ggml_type qtypes[] = { GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, GGML_TYPE_Q8_0 };
const int q_ne01[] = { 1, 2, 3, 8 };
const int q_ne11[] = { 512, 4096, 8192 };
for (ggml_type qt : qtypes) {
for (auto s : shapes) {
for (int ne : fa_vec_legal_ne(s.dk, s.dv)) {
for (int Q : { 1, 2, 4 }) {
for (bool sinks : { false, true }) {
for (int ne01 : q_ne01) {
for (int ne11 : q_ne11) {
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, qt, qt);
auto st = tc.eval(backend_metal, backend_cpu, "FLASH_ATTN_EXT", nullptr);
clear_ov();
if (st == test_status_t::FAIL) {
printf("FAIL(quant) type=%s dk=%d dv=%d Q=%d ne=%d ne01=%d ne11=%d sinks=%d\n",
ggml_type_name(qt), s.dk, s.dv, Q, ne, ne01, ne11, (int) sinks);
ok = false; n_fail++;
}
n_run++;
}
}
}
}
}
}
}
printf("fa_vec tune-check: %d cases run, %d failed\n", n_run, n_fail);
return ok;
}
// A prebuilt FA op graph for one (dk,dv,ne01,ne11) cell. The op is replicated n_runs
// times so a single graph_compute amortizes dispatch/sync overhead. The graph is reused
// across (Q,NE) overrides (the override only changes the pipeline picked at encode time),
// which avoids re-allocating/re-initializing the (large) K/V tensors per candidate.
struct fa_perf_cell {
ggml_context_ptr ctx;
ggml_backend_buffer_ptr buf;
ggml_cgraph * gf = nullptr;
int n_runs = 0;
bool ok = false;
};
static fa_perf_cell fa_build_perf_cell(ggml_backend_t backend, int dk, int dv, int ne01, int ne11,
ggml_type type_kv = GGML_TYPE_F16) {
fa_perf_cell cell;
// GQA shape (nr23=[8,1]) matching real spec-decode / verify workloads: enough query
// heads to keep the GPU busy so the Q>1 K/V-reuse benefit is visible (and comparable
// to the documented #23114 numbers). nh here is the number of KV heads.
test_flash_attn_ext tc(dk, dv, /*nh=*/4, { 8, 1 }, /*kv=*/ne11, /*nb=*/ne01,
/*mask=*/true, /*sinks=*/false, 0.0f, 0.0f, GGML_PREC_F32,
type_kv, type_kv);
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 = tc.build_graph(cell.ctx.get());
if (!ggml_backend_supports_op(backend, out)) {
return cell;
}
cell.buf.reset(ggml_backend_alloc_ctx_tensors(cell.ctx.get(), backend));
if (cell.buf == NULL) {
return cell;
}
tc.initialize_tensors(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;
if (tc.op_flops(out) > 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 / tc.op_flops(out)));
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);
}
cell.ok = true;
return cell;
}
// Median per-op GPU time (us) for the currently-set override over the prebuilt cell graph.
static double time_fa_cell_median(ggml_backend_t backend, const fa_perf_cell & cell, int reps) {
if (!cell.ok) {
return -1.0;
}
ggml_backend_graph_compute(backend, cell.gf); // warmup (compiles the pipeline for the override)
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;
}
// Perf sweep over the (Q,NE) grid, emitting pasteable fa_vec_tuned_table rows.
// nsg/nwg are left to the ops.cpp adaptive heuristic (not part of the table).
static bool run_fa_vec_tune_perf(ggml_backend_t backend_metal) {
auto * reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend_metal));
auto set_ov = (set_fa_vec_override_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_metal_tuning_set_fa_vec_override");
auto clr_ov = (clear_fa_vec_override_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_metal_tuning_clear_fa_vec_override");
// Share the runtime's bucketers + baseline NE via read-only proc bridges, so the emitted
// (ne11_b, ne01_b) keys match fa_vec_pick and can't silently desync from FA_VEC_*_BUCKETS.
auto ne11_bucket = (fa_vec_bucket_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_metal_tuning_fa_vec_ne11_bucket");
auto ne01_bucket = (fa_vec_bucket_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_metal_tuning_fa_vec_ne01_bucket");
auto baseline_ne = (fa_vec_baseline_ne_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_metal_tuning_fa_vec_baseline_ne");
if (!set_ov || !clr_ov || !ne11_bucket || !ne01_bucket || !baseline_ne) {
printf("metal fa_vec tuning procs unavailable\n");
return false;
}
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 } };
const int ne11_rep[] = { 512, 2048, 8192, 32768 }; // ne11 bucket representatives
const int ne01_rep[] = { 1, 2, 3, 4, 5, 6, 7, 8, 16 }; // PR grid: point-bucket reps (1-4) + tail mod-4 cycle (5-8) + large anchor (16); vec serves ne01<20
const int REPS = 7; // odd -> exact median
printf("# fa_vec perf sweep — replace GGML_METAL_DEVICE_M4_MAX with this machine's device\n");
printf("# [1] per-bucket timings (us, winner *): '=> cfg Nx' beats baseline, else '=> baseline'\n");
printf("# [2] a pasteable fa_vec_tuned_table block (domain defaults + exceptions) is printed after [1]\n");
struct cand_t { int Q, NE; double t; }; // one timed (Q,NE) candidate
struct pt_t { int dk, dv, ne11, ne01; // one swept grid point with its candidate times
std::vector<cand_t> cs; double base_t; };
// group_split compression knobs (see ggml-metal-tuning.h for the row / lookup semantics)
const double TUNE_TAU = 0.05; // max POINTWISE regret (worst point in a bucket) to ride a domain
// default instead of emitting the bucket's own exception row
const double TUNE_THETA = 1.05; // min AGGREGATE bucket speedup vs baseline to tune the bucket at all
auto type_token = [](ggml_type t) -> const char * {
switch (t) {
case GGML_TYPE_Q4_0: return "GGML_TYPE_Q4_0";
case GGML_TYPE_Q4_1: return "GGML_TYPE_Q4_1";
case GGML_TYPE_Q5_0: return "GGML_TYPE_Q5_0";
case GGML_TYPE_Q5_1: return "GGML_TYPE_Q5_1";
case GGML_TYPE_Q8_0: return "GGML_TYPE_Q8_0";
default: return "GGML_TYPE_F16";
}
};
const ggml_type types[] = { GGML_TYPE_F16, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1,
GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, GGML_TYPE_Q8_0 };
for (ggml_type type_kv : types) {
printf("\n### dtype=%s\n", ggml_type_name(type_kv));
std::vector<pt_t> pts; // one per swept (shape, ne11, ne01); bucketed + compressed below
for (auto s : shapes) {
const int base_ne = baseline_ne(s.dk, s.dv);
const std::vector<int> legal = fa_vec_legal_ne(s.dk, s.dv);
for (int ne11 : ne11_rep) {
for (int ne01 : ne01_rep) {
fa_perf_cell cell = fa_build_perf_cell(backend_metal, s.dk, s.dv, ne01, ne11, type_kv);
if (!cell.ok) {
continue;
}
std::vector<cand_t> cs;
for (int ne : legal) {
for (int Q : { 1, 2, 4 }) {
cs.push_back({ Q, ne, 0.0 });
}
}
// randomize config order to decorrelate thermal throttling across the cell
std::shuffle(cs.begin(), cs.end(), std::mt19937(1234));
double anchor = 0.0; // periodic baseline re-measure -> throttling detection
for (size_t i = 0; i < cs.size(); ++i) {
set_ov(cs[i].Q, cs[i].NE);
cs[i].t = time_fa_cell_median(backend_metal, cell, REPS);
clr_ov();
if (i % 4 == 0) {
const double a = time_fa_cell_median(backend_metal, cell, REPS);
if (anchor > 0.0) {
double drift = (a - anchor) / anchor;
if (drift < 0) {
drift = -drift;
}
if (drift > 0.10) {
printf("# WARN throttling? anchor drift %.1f%% dk=%d ne11=%d\n", 100.0 * drift, s.dk, ne11);
}
}
anchor = (anchor > 0.0) ? std::min(anchor, a) : a;
}
}
std::sort(cs.begin(), cs.end(), [](const cand_t & a, const cand_t & b) {
return a.Q != b.Q ? a.Q < b.Q : a.NE < b.NE;
});
cand_t best = cs[0];
for (const auto & c : cs) {
if (c.t > 0.0 && (best.t <= 0.0 || c.t < best.t)) {
best = c;
}
}
double base_t = 0.0; // the swept (Q=1, base_ne) candidate == baseline kernel
for (const auto & c : cs) {
if (c.Q == 1 && c.NE == base_ne) { base_t = c.t; break; }
}
const bool keep = best.t > 0.0 && base_t > 0.0 && best.t < base_t * 0.98;
printf("# dtype=%s dk=%d dv=%d ne11=%d ne01=%d:", ggml_type_name(type_kv), s.dk, s.dv, ne11, ne01);
for (const auto & c : cs) {
printf(" Q%dNE%d=%.1f%s", c.Q, c.NE, c.t, (c.Q == best.Q && c.NE == best.NE) ? "*" : "");
}
if (keep) {
printf(" => Q%d,NE%d %.2fx\n", best.Q, best.NE, base_t / best.t);
} else {
printf(" => baseline\n");
}
pts.push_back({ s.dk, s.dv, ne11, ne01, cs, base_t });
}
}
}
// [2] Compress into pasteable rows. Per (dk,dv) and ne01 domain {decode==1, batch>=2}, emit
// one ne11-collapsed default cfg (fewest rows, ne11_b=-1) plus a per-bucket exception wherever
// the default's pointwise regret vs the bucket's min-max-regret target, or its aggregate
// slowdown vs baseline, exceeds TUNE_TAU. The target still carries an intrinsic per-point regret
// from lumping ne01 into one tail bucket — TUNE_TAU bounds resolved-vs-target, not vs-oracle.
std::vector<std::string> rows_out;
char rbuf[192];
for (auto s : shapes) {
const int base_ne = baseline_ne(s.dk, s.dv);
std::vector<cand_t> cfgs; // candidate list (identical for every grid point of this shape)
int base_i = 0;
for (int ne : fa_vec_legal_ne(s.dk, s.dv)) {
for (int Q : { 1, 2, 4 }) {
if (Q == 1 && ne == base_ne) { base_i = (int) cfgs.size(); }
cfgs.push_back({ Q, ne, 0.0 });
}
}
auto cfg_time = [&](const pt_t & p, int i) {
for (const auto & c : p.cs) { if (c.Q == cfgs[i].Q && c.NE == cfgs[i].NE) { return c.t; } }
return 0.0;
};
// Bucket the grid points using the runtime's bucketers (via proc), so keys match fa_vec_pick.
// Short-KV points (ne11 bucket 0) are dropped — the runtime serves those from baseline. Each
// kept bucket picks a min-max-regret target (gated by the aggregate TUNE_THETA benefit gate).
struct bkt_t { int b11, b01, Ti; std::vector<double> agg; double base_agg;
std::vector<const pt_t *> bp; }; // bp kept for the pointwise ride test below
std::set<std::pair<int, int>> seen;
for (const auto & p : pts) {
if (p.dk != s.dk || p.dv != s.dv) { continue; }
const int b11 = ne11_bucket(p.ne11);
if (b11 == 0) { continue; }
seen.insert({ b11, ne01_bucket(p.ne01) });
}
std::vector<bkt_t> bks;
for (const auto & bb : seen) {
const int b11 = bb.first, b01 = bb.second;
std::vector<const pt_t *> bp;
for (const auto & p : pts) {
if (p.dk == s.dk && p.dv == s.dv && ne11_bucket(p.ne11) == b11 && ne01_bucket(p.ne01) == b01) {
bp.push_back(&p);
}
}
std::vector<double> agg(cfgs.size(), 0.0), worst(cfgs.size(), 0.0);
for (const auto * p : bp) {
double bestt = 0.0;
for (size_t i = 0; i < cfgs.size(); ++i) {
double t = cfg_time(*p, (int) i);
if (t > 0.0 && (bestt == 0.0 || t < bestt)) { bestt = t; }
}
for (size_t i = 0; i < cfgs.size(); ++i) {
double t = cfg_time(*p, (int) i);
agg[i] += t;
if (t > 0.0 && bestt > 0.0) { worst[i] = std::max(worst[i], t / bestt); }
}
}
int robust = 0;
for (size_t i = 1; i < cfgs.size(); ++i) {
if (worst[i] < worst[robust] ||
(worst[i] == worst[robust] && (cfgs[i].Q < cfgs[robust].Q ||
(cfgs[i].Q == cfgs[robust].Q && cfgs[i].NE < cfgs[robust].NE)))) {
robust = (int) i;
}
}
const bool tune = robust != base_i && agg[base_i] / agg[robust] >= TUNE_THETA;
bks.push_back({ b11, b01, tune ? robust : base_i, agg, agg[base_i], bp });
}
// Pointwise regret of default cfg d vs the bucket target Ti. Per-point (not aggregate):
// a ratio-of-sums lets a default that wins on aligned ne01 (8, 16) hide a large penalty
// on a misaligned point (ne01=5), so the ride/exception decision must be per point.
auto reg_pointwise = [&](const bkt_t * b, int d) {
double r = 0.0;
for (const auto * p : b->bp) {
const double td = cfg_time(*p, d), tT = cfg_time(*p, 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 < cfgs.size(); ++d) {
int rows = ((int) d != base_i) ? 1 : 0; double tot = 0.0;
for (const auto * b : db) {
const double reg = reg_pointwise(b, (int) d); // pointwise vs bucket target
const double slow = b->agg[d] / b->base_agg - 1.0; // aggregate vs baseline (safety net)
if (reg > TUNE_TAU || slow > TUNE_TAU) { rows++; tot += b->agg[b->Ti]; }
else { tot += b->agg[d]; }
}
const bool better = bestD < 0 || rows < bestRows ||
(rows == bestRows && (tot < bestTot ||
(tot == bestTot && (cfgs[d].Q < cfgs[bestD].Q ||
(cfgs[d].Q == cfgs[bestD].Q && cfgs[d].NE < cfgs[bestD].NE)))));
if (better) { bestD = (int) d; bestRows = rows; bestTot = tot; }
}
const int dom_id = (dom == 0) ? 0 : 1; // FA_VEC_DOMAIN_DECODE / FA_VEC_DOMAIN_BATCH
if (bestD != base_i) {
snprintf(rbuf, sizeof(rbuf), " { { GGML_METAL_DEVICE_M4_MAX, %s, %d, %d, -1, %d }, { %d, %d } },",
type_token(type_kv), s.dk, s.dv, dom_id, cfgs[bestD].Q, cfgs[bestD].NE);
rows_out.emplace_back(rbuf);
}
for (const auto * b : db) {
const double reg = reg_pointwise(b, bestD); // pointwise vs bucket target
const double slow = b->agg[bestD] / b->base_agg - 1.0; // aggregate vs baseline
if (reg <= TUNE_TAU && slow <= TUNE_TAU) { continue; } // rides the default / baseline
snprintf(rbuf, sizeof(rbuf), " { { GGML_METAL_DEVICE_M4_MAX, %s, %d, %d, %d, %d }, { %d, %d } },",
type_token(type_kv), s.dk, s.dv, b->b11, b->b01, cfgs[b->Ti].Q, cfgs[b->Ti].NE);
rows_out.emplace_back(rbuf);
}
}
}
printf("\n // ---- %s: %zu rows ----\n", ggml_type_name(type_kv), rows_out.size());
for (const auto & r : rows_out) { printf("%s\n", r.c_str()); }
} // for type_kv
return true;
}
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) {
printer * output_printer, const char * test_file_path, int parallel_workers, bool tune_perf = false) {
auto filter_test_cases = [](std::vector<std::unique_ptr<test_case>> & test_cases, const char * params_filter) {
if (params_filter == nullptr) {
return;
@@ -10142,6 +10617,9 @@ static bool test_backend(ggml_backend_t backend, ggml_backend_dev_t dev, test_mo
case MODE_PERF:
test_cases = make_test_cases_perf();
break;
case MODE_TUNE:
// MODE_TUNE routes to its own dispatch below; no generic test_cases needed
break;
}
} else {
test_cases = make_test_cases_from_file(test_file_path);
@@ -10277,6 +10755,29 @@ static bool test_backend(ggml_backend_t backend, ggml_backend_dev_t dev, test_mo
return true;
}
if (mode == MODE_TUNE) {
// self-create a CPU backend with the reference implementation as golden.
// (backend_cpu in the MODE_TEST block above is out of scope here.)
ggml_backend_t backend_cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, NULL);
GGML_ASSERT(backend_cpu != NULL);
{
using ggml_backend_cpu_set_use_ref_t = void (*)(ggml_backend_t, bool);
auto * cpu_reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend_cpu));
auto * set_use_ref = (ggml_backend_cpu_set_use_ref_t) ggml_backend_reg_get_proc_address(cpu_reg, "ggml_backend_cpu_set_use_ref");
if (set_use_ref) {
set_use_ref(backend_cpu, true);
}
}
// backend here is the MODE_TUNE-selected metal backend (-b MTL0)
const bool ok = tune_perf
? run_fa_vec_tune_perf(backend)
: (run_fa_vec_drift_guard(backend, backend_cpu) && run_fa_vec_tune_check(backend, backend_cpu));
ggml_backend_free(backend_cpu);
return ok;
}
if (mode == MODE_SUPPORT) {
// Filter out fusion cases
test_cases.erase(
@@ -10402,6 +10903,7 @@ static void usage(char ** argv) {
printf(" - grad (compare gradients from backpropagation with method of finite differences)\n");
printf(" - perf (performance evaluation)\n");
printf(" - support (probe backend operation support)\n");
printf(" - tune (FA vec (Q,NE) numerical correctness check vs CPU reference; --tune-perf for the perf sweep)\n");
printf(" op names for -o are as given by ggml_op_desc() (e.g. ADD, MUL_MAT, etc),\n");
printf(" optionally including the full test case string (e.g. \"ADD(type=f16,ne=[1,1,8,1],nr=[1,1,1,1],nf=1)\")\n");
printf(" --output specifies output format (default: console, options: console, sql, csv)\n");
@@ -10419,6 +10921,7 @@ int main(int argc, char ** argv) {
const char * params_filter = nullptr;
const char * test_file_path = nullptr;
int parallel_workers = 1;
bool tune_perf = false;
for (int i = 1; i < argc; i++) {
if (strcmp(argv[i], "test") == 0) {
@@ -10429,6 +10932,10 @@ int main(int argc, char ** argv) {
mode = MODE_GRAD;
} else if (strcmp(argv[i], "support") == 0) {
mode = MODE_SUPPORT;
} else if (strcmp(argv[i], "tune") == 0) {
mode = MODE_TUNE;
} else if (strcmp(argv[i], "--tune-perf") == 0) {
tune_perf = true;
} else if (strcmp(argv[i], "-o") == 0) {
if (i + 1 < argc) {
op_names_filter = argv[++i];
@@ -10536,7 +11043,7 @@ int main(int argc, char ** argv) {
false, "", ggml_backend_dev_description(dev),
total / 1024 / 1024, free / 1024 / 1024, true));
bool ok = test_backend(backend.get(), dev, mode, op_names_filter, params_filter, output_printer.get(), test_file_path, parallel_workers);
bool ok = test_backend(backend.get(), dev, mode, op_names_filter, params_filter, output_printer.get(), test_file_path, parallel_workers, tune_perf);
if (ok) {
n_ok++;