metal : per-device tuned (Q, NE) for flash-attn vec (#26570)
* 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 * add device id -> token lookup for the offline tuning tool * add ggml-metal-tuning skeleton * add op-agnostic perf cell + median timing for the tuner * add FA-vec graph build + tensor init to the tuner * tools : add FA-vec (Q,NE) sweep, compression and table emit * cool down and re-measure the dirty window on thermal drift * test-backend-ops : replace the FA vec tune mode with a bounded (Q,NE) slice * tools : document the Metal tuner, point the table comment at it * abort on unknown KV type, single-source fa_vec_legal_ne * cleanup * honor -o in the FA vec (Q,NE) slice * retune FA-vec (Q, NE) under a pointwise no-harm gate * cont : add fa-vec tunings for M1 Pro, M2 Ultra, M5 Max --------- Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
This commit is contained in:
co-authored by
Georgi Gerganov
parent
b615f5b4bd
commit
f280b26983
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user