char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

YiChen Lvmetal : per-device tuned (Q, NE) for flash-attn vec (#26570)f280b2698

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
2.7 KiB77 linesraw
1#pragma once
2
3#include "ggml-metal-device.h"  // enum ggml_metal_device_id
4#include "ggml.h"
5
6#include <cstdint>
7#include <vector>
8
9namespace ggml_metal_tuning {
10
11// FA vec selection buckets. ne01 (query rows) splits decode (==1) from batch (>=2), the
12// batch side refined into {2,3,4,5}: Q>1 reuses one K/V load across rows, so it only pays
13// off once ne01 aligns with Q. ne11 (KV length) is bucketed too, as the Q>1 crossover is
14// head-size dependent (small dk crosses late, large dk wins even at short KV).
15constexpr int FA_VEC_NE11_BUCKETS[] = { 1024, 4096, 16384 };
16constexpr int FA_VEC_NE01_BUCKETS[] = { 2, 3, 4, 5 };
17
18int fa_vec_ne11_bucket(int64_t ne11);
19int fa_vec_ne01_bucket(int64_t ne01);
20
21// NE baked into each (dk,dv) baseline instantiation in kernels/fa.metal.
22// Hand-maintained mirror; keep in sync with those instantiations.
23// The Metal test slice covers every legal config for dk=128 and dk=576.
24int fa_vec_baseline_ne(int dk, int dv);
25
26// Tuned table has two row kinds. Exact rows key a (ne11_b, ne01_b) bucket. Default rows
27// collapse ne11 over one ne01 domain: ne11_b == FA_VEC_NE11_DEFAULT and ne01_b holds the
28// domain. fa_vec_pick tries exact bucket -> domain default -> baseline; short KV
29// (ne11 < FA_VEC_NE11_BUCKETS[0]) always uses baseline.
30constexpr int8_t FA_VEC_NE11_DEFAULT  = -1;
31constexpr int8_t FA_VEC_DOMAIN_DECODE = 0;  // ne01 == 1
32constexpr int8_t FA_VEC_DOMAIN_BATCH  = 1;  // ne01 >= 2
33
34struct fa_vec_key_t {
35    int8_t  device_id;
36    int8_t  dtype;
37    int16_t dk;
38    int16_t dv;
39    int8_t  ne11_b;
40    int8_t  ne01_b;
41};
42
43static_assert(sizeof(fa_vec_key_t) == 8, "fa_vec_key_t must be tightly packed for memcmp");
44
45struct fa_vec_cfg_t {
46    int8_t Q;
47    int8_t NE;
48};
49
50struct fa_vec_entry_t {
51    fa_vec_key_t key;
52    fa_vec_cfg_t cfg;
53};
54
55// legal NE values for a (dk,dv): NL = 32/NE, require (dk/4)%NL==0 && (dv/4)%NL==0.
56// single source shared by the offline tuner and test-backend-ops.
57inline std::vector<int> fa_vec_legal_ne(int dk, int dv) {
58    std::vector<int> r;
59    for (int ne : { 1, 2, 4 }) {
60        const int nl = 32 / ne;
61        if ((dk / 4) % nl == 0 && (dv / 4) % nl == 0) {
62            r.push_back(ne);
63        }
64    }
65    return r;
66}
67
68// test/tune-only override; when set, fa_vec_pick returns it directly.
69void         fa_vec_set_override(fa_vec_cfg_t cfg);
70void         fa_vec_clear_override();
71fa_vec_cfg_t fa_vec_baseline_cfg(int dk, int dv);
72
73// device_id selects a per-SKU row; on a miss, gpu_family (0 if unknown) maps to a representative
74// SKU and the table is retried. No match -> baseline.
75fa_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);
76
77}  // namespace ggml_metal_tuning