diff --git a/tools/tuning/fa-vec.cpp b/tools/tuning/fa-vec.cpp index e21e17a4d0..38060576ab 100644 --- a/tools/tuning/fa-vec.cpp +++ b/tools/tuning/fa-vec.cpp @@ -1,9 +1,194 @@ #include "fa-vec.h" +#include "bench.h" +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include #include +#include +#include +#include +#include +#include + +// 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; +} + +// mirrors init_tensor_uniform: uniform f32 data, quantized in place for quantized types +static void fa_init_uniform(ggml_tensor * t, std::mt19937 & rng, float min, float max) { + const size_t nels = ggml_nelements(t); + + std::vector data(nels); + std::uniform_real_distribution 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 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 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 data_f32(size_t(ne0)*ne1*ne2*ne3); + std::vector data_f16(size_t(ne0)*ne1*ne2*ne3); + + std::uniform_real_distribution 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)); +} + +// per-cell deterministic seed: the shape decides it, so a cell is reproducible +// regardless of what else the sweep visited before it +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); + } + } +} + +// legal NE for a (dk,dv): NL = 32/NE, require (dk/4)%NL==0 && (dv/4)%NL==0 +static std::vector fa_legal_ne(int dk, int dv) { + std::vector 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; +} bool tuner_fa_vec_run(ggml_backend_t backend, ggml_backend_dev_t dev, const tuner_opts & opts) { - fprintf(stderr, "fa-vec tuner: not implemented yet\n"); + fprintf(stderr, "fa-vec tuner: sweep not implemented yet\n"); (void) backend; (void) dev;