mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-08 08:12:16 -10:00
Merge f-mac-kern 272bbbbd
This commit is contained in:
Niko1221 committed
commit
f0609f40d4
8 files changed
+269
-11
No files matched your search
@@ -321,3 +321,165 @@ template [[host_name("gdn_scanr_2")]] kernel decltype(gdn_scan_r<2, 16>) gdn_sca
|
||||
template [[host_name("gdn_scanr_4")]] kernel decltype(gdn_scan_r<4, 16>) gdn_scan_r<4, 16>;
|
||||
template [[host_name("gdn_scanr_8")]] kernel decltype(gdn_scan_r<8, 16>) gdn_scan_r<8, 16>;
|
||||
template [[host_name("gdn_scanr_4_8")]] kernel decltype(gdn_scan_r<4, 8>) gdn_scan_r<4, 8>;
|
||||
|
||||
// ---- the chunked form (chunks of GC = 32 tokens), the WY / UT form of the gated delta rule. With G_t the product
|
||||
// of the chunk's gates up to t (log space: lg_t) and S0 the state entering the chunk:
|
||||
// (I + A) D = B, A_ts = beta_t (k_t . k_s) G_t / G_s (s < t), B = [beta v | beta G k] -> D = [U | W]
|
||||
// delta = U - W S0, o_t = G_t S0^T q_t + sum_{s <= t} (q_t . k_s) G_t / G_s delta_s,
|
||||
// S_L = G_L S0 + sum_s G_L / G_s k_s delta_s^T
|
||||
// (the sequential recurrence's algebra; the rounding differs). Three kernels:
|
||||
// gdn_cprep per (head, chunk): lg, A, the forward substitution for U and W (a thread per column of B)
|
||||
// gdn_cstate per (head, 32 state columns): the chunks in order - delta, the inter-chunk output, the state update
|
||||
// gdn_cout per (head, chunk): o += the intra-chunk part, times 1/sqrt(128)
|
||||
constant constexpr int GC = 32;
|
||||
|
||||
// grid (h_v, chunks), 256 threads. U, W: [h][chunk][GC][128]; lgo: [h][chunk][GC]
|
||||
kernel void gdn_cprep(device const float* qn [[buffer(0)]], device const float* kn [[buffer(1)]],
|
||||
device const float* vc [[buffer(2)]], device const float2* gb [[buffer(3)]],
|
||||
device float* U [[buffer(4)]], device float* W [[buffer(5)]], device float* lgo [[buffer(6)]],
|
||||
constant GdnPArgs& a [[buffer(7)]],
|
||||
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) {
|
||||
constexpr int S = 128;
|
||||
threadgroup float A[GC][GC];
|
||||
threadgroup float lg[GC], bt[GC];
|
||||
const int h = int(tg.x), c = int(tg.y), hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok);
|
||||
const int NC = (T + GC - 1) / GC, t0 = c * GC, nb = min(GC, T - t0);
|
||||
if (tid == 0) {
|
||||
float s = 0.0f;
|
||||
for (int t = 0; t < GC; ++t) {
|
||||
if (t < nb) s += max(log(gb[(t0 + t) * hv + h].x), -80.0f);
|
||||
lg[t] = s;
|
||||
bt[t] = t < nb ? gb[(t0 + t) * hv + h].y : 0.0f;
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (int idx = tid; idx < GC * GC; idx += 256) {
|
||||
const int t = idx / GC, s = idx % GC;
|
||||
float v = 0.0f;
|
||||
if (s < t && t < nb) {
|
||||
device const float4* kt = (device const float4*)(kn + ((t0 + t) * hkn + hk) * S);
|
||||
device const float4* ks = (device const float4*)(kn + ((t0 + s) * hkn + hk) * S);
|
||||
float d = 0.0f;
|
||||
for (int i = 0; i < S / 4; ++i) d += dot(kt[i], ks[i]);
|
||||
v = bt[t] * d * exp(lg[t] - lg[s]);
|
||||
}
|
||||
A[t][s] = v;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
// forward substitution, column j of B a thread (threads 0..127: beta v; 128..255: beta G k)
|
||||
const int j = tid & 127;
|
||||
const bool isw = tid >= 128;
|
||||
float x[GC];
|
||||
for (int t = 0; t < GC; ++t) {
|
||||
float b = 0.0f;
|
||||
if (t < nb) b = isw ? bt[t] * exp(lg[t]) * kn[((t0 + t) * hkn + hk) * S + j] : bt[t] * vc[((t0 + t) * hv + h) * S + j];
|
||||
for (int s = 0; s < t; ++s) b = fma(-A[t][s], x[s], b);
|
||||
x[t] = b;
|
||||
}
|
||||
device float* o = (isw ? W : U) + (ulong(h) * NC + c) * GC * S + j;
|
||||
for (int t = 0; t < GC; ++t) o[t * S] = x[t];
|
||||
if (tid < GC) lgo[(h * NC + c) * GC + tid] = lg[tid];
|
||||
}
|
||||
|
||||
// grid (h_v * 4), 256 threads = 8 row blocks of 16 x 32 columns; the state's 128 x 32 block in registers (16 a thread)
|
||||
kernel void gdn_cstate(device float* state [[buffer(0)]], device const float* qn [[buffer(1)]],
|
||||
device const float* kn [[buffer(2)]], device const float* U [[buffer(3)]],
|
||||
device const float* W [[buffer(4)]], device const float* lgo [[buffer(5)]],
|
||||
device float* delta [[buffer(6)]], device float* o [[buffer(7)]], constant GdnPArgs& a [[buffer(8)]],
|
||||
uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
|
||||
ushort lane [[thread_index_in_simdgroup]]) {
|
||||
constexpr int S = 128, RB = 16;
|
||||
threadgroup float red[8][16][32];
|
||||
threadgroup float dl[GC][32];
|
||||
const int h = int(tg.x) / 4, hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok);
|
||||
const int NC = (T + GC - 1) / GC;
|
||||
const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * RB;
|
||||
float st[RB];
|
||||
for (int r = 0; r < RB; ++r) st[r] = state[((i0 + r) * hv + h) * S + j];
|
||||
for (int c = 0; c < NC; ++c) {
|
||||
const int t0 = c * GC, nb = min(GC, T - t0);
|
||||
device const float* Uc = U + (ulong(h) * NC + c) * GC * S;
|
||||
device const float* Wc = W + (ulong(h) * NC + c) * GC * S;
|
||||
device const float* lg = lgo + (h * NC + c) * GC;
|
||||
// pass 0: delta = U - W S0; pass 1: o_inter = G_t S0^T q_t; 16 tokens at a time
|
||||
for (int pass = 0; pass < 2; ++pass) {
|
||||
for (int hf = 0; hf < 2; ++hf) {
|
||||
float p[16];
|
||||
for (int u = 0; u < 16; ++u) {
|
||||
const int t = hf * 16 + u;
|
||||
float s = 0.0f;
|
||||
if (t < nb) {
|
||||
device const float* v = pass == 0 ? Wc + t * S + i0 : qn + ((t0 + t) * hkn + hk) * S + i0;
|
||||
for (int r = 0; r < RB; ++r) s = fma(v[r], st[r], s);
|
||||
}
|
||||
p[u] = s;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (int u = 0; u < 16; ++u) red[sg][u][lane] = p[u];
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
for (int u = sg * 2; u < sg * 2 + 2; ++u) {
|
||||
const int t = hf * 16 + u;
|
||||
float s = 0.0f;
|
||||
for (int b = 0; b < 8; ++b) s += red[b][u][lane];
|
||||
if (t < nb) {
|
||||
if (pass == 0) {
|
||||
const float d = Uc[t * S + j] - s;
|
||||
dl[t][lane] = d;
|
||||
delta[((ulong(h) * NC + c) * GC + t) * S + j] = d;
|
||||
} else {
|
||||
o[((t0 + t) * hv + h) * S + j] = exp(lg[t]) * s;
|
||||
}
|
||||
} else if (pass == 0) {
|
||||
dl[t][lane] = 0.0f;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
// S = G_L S0 + sum_s G_L / G_s k_s delta_s^T
|
||||
const float lgl = lg[nb - 1], gl = exp(lgl);
|
||||
for (int r = 0; r < RB; ++r) st[r] *= gl;
|
||||
for (int s = 0; s < nb; ++s) {
|
||||
const float f = exp(lgl - lg[s]) * dl[s][lane];
|
||||
device const float* ks = kn + ((t0 + s) * hkn + hk) * S + i0;
|
||||
for (int r = 0; r < RB; ++r) st[r] = fma(ks[r], f, st[r]);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
for (int r = 0; r < RB; ++r) state[((i0 + r) * hv + h) * S + j] = st[r];
|
||||
}
|
||||
|
||||
// grid (h_v, chunks), 128 threads (column j): o_t = (o_inter_t + sum_{s <= t} (q_t . k_s) G_t / G_s delta_s) / sqrt(128)
|
||||
kernel void gdn_cout(device const float* qn [[buffer(0)]], device const float* kn [[buffer(1)]],
|
||||
device const float* delta [[buffer(2)]], device const float* lgo [[buffer(3)]],
|
||||
device float* o [[buffer(4)]], constant GdnPArgs& a [[buffer(5)]],
|
||||
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) {
|
||||
constexpr int S = 128;
|
||||
threadgroup float M[GC][GC];
|
||||
const int h = int(tg.x), c = int(tg.y), hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok);
|
||||
const int NC = (T + GC - 1) / GC, t0 = c * GC, nb = min(GC, T - t0);
|
||||
device const float* lg = lgo + (h * NC + c) * GC;
|
||||
for (int idx = tid; idx < GC * GC; idx += 128) {
|
||||
const int t = idx / GC, s = idx % GC;
|
||||
float v = 0.0f;
|
||||
if (s <= t && t < nb) {
|
||||
device const float4* qt = (device const float4*)(qn + ((t0 + t) * hkn + hk) * S);
|
||||
device const float4* ks = (device const float4*)(kn + ((t0 + s) * hkn + hk) * S);
|
||||
float d = 0.0f;
|
||||
for (int i = 0; i < S / 4; ++i) d += dot(qt[i], ks[i]);
|
||||
v = d * exp(lg[t] - lg[s]);
|
||||
}
|
||||
M[t][s] = v;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
const int j = tid;
|
||||
device const float* dc = delta + (ulong(h) * NC + c) * GC * S + j;
|
||||
float dv[GC];
|
||||
for (int s = 0; s < GC; ++s) dv[s] = s < nb ? dc[s * S] : 0.0f;
|
||||
const float scale = 1.0f / sqrt(float(S));
|
||||
for (int t = 0; t < nb; ++t) {
|
||||
float s = o[((t0 + t) * hv + h) * S + j];
|
||||
for (int u = 0; u <= t; ++u) s = fma(M[t][u], dv[u], s);
|
||||
o[((t0 + t) * hv + h) * S + j] = s * scale;
|
||||
}
|
||||
}
|
||||
@@ -45,3 +45,19 @@ kernel void moe_combine_sh(device const float* ye [[buffer(0)]], device const in
|
||||
const float gs = 1.0f / (1.0f + exp(-sg[t]));
|
||||
bo[gid] = fma(gs, ysh[gid], s);
|
||||
}
|
||||
|
||||
// slc_touch: reads n16 16-byte words (a weight range the next kernels will need) so they sit in the SLC; the xor is
|
||||
// written only if it equals a value it never takes (keeps the loads alive). grid (ceil(n16 / 1024)), 256 threads,
|
||||
// 4 words a thread.
|
||||
kernel void slc_touch(device const uint4* src [[buffer(0)]], device uint* sink [[buffer(1)]], constant uint& n16 [[buffer(2)]],
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
uint s = 0;
|
||||
for (uint k = 0; k < 4; ++k) {
|
||||
const uint j = (gid >> 8) * 1024 + k * 256 + (gid & 255);
|
||||
if (j < n16) {
|
||||
const uint4 v = src[j];
|
||||
s ^= v.x ^ v.y ^ v.z ^ v.w;
|
||||
}
|
||||
}
|
||||
if (s == 0x9e3779b9u) sink[0] = s;
|
||||
}
|
||||
@@ -82,3 +82,4 @@ kernel void coop_layout(device int* out [[buffer(0)]], ushort lane [[thread_inde
|
||||
template [[host_name("coop_layout_16_32_16")]] kernel decltype(coop_layout<16, 32, 16>) coop_layout<16, 32, 16>;
|
||||
template [[host_name("coop_layout_32_32_16")]] kernel decltype(coop_layout<32, 32, 16>) coop_layout<32, 32, 16>;
|
||||
template [[host_name("coop_layout_16_64_32")]] kernel decltype(coop_layout<16, 64, 32>) coop_layout<16, 64, 32>;
|
||||
template [[host_name("coop_layout_16_32_32")]] kernel decltype(coop_layout<16, 32, 32>) coop_layout<16, 32, 32>;
|
||||
@@ -406,6 +406,7 @@ static inline half4 w4_signed(uint g, uint s, float sc) {
|
||||
return half4(q * sc);
|
||||
}
|
||||
template <> struct W4<T_Q2_0> { // 64 values a block: fp16 d, 16 bytes of 2-bit codes (value v at byte v / 4)
|
||||
// fp32 math then fp16 (an fp16-math decode was measured slower: 4.2 vs 3.5 ms a layer)
|
||||
static inline half4 run(device const uchar* row, int k, threadgroup const uint*) {
|
||||
device const uchar* b = row + (k >> 6) * 18;
|
||||
const float d = float(*(device const half*)b);
|
||||
@@ -495,6 +496,7 @@ template <int F> struct W8 {
|
||||
b = W4<F>::run(row, k + 16, tab);
|
||||
}
|
||||
};
|
||||
// (a W8<T_Q2_0> sharing the block scale was measured slower than the two W4 of the generic one: 5.5 vs 3.5-4.2 ms)
|
||||
template <> struct W8<T_IQ3_XXS> {
|
||||
static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) {
|
||||
device const uchar* bl = row + (k >> 8) * 98;
|
||||
@@ -612,8 +614,11 @@ kernel void mm_gu3(device const uchar* W [[buffer(0)]], device const half* XS [[
|
||||
device const uchar* rgw = g0 + ulong(n) * a.gu_row;
|
||||
device const uchar* ruw = u0 + ulong(n) * a.gu_row;
|
||||
half4 wg, wu, wg2, wu2;
|
||||
if (TK == 32 && iy[0] == ix[0] + 16 && iy[1] == ix[1] && (ix[0] & 31) < 16) {
|
||||
// the lane's two runs of this row sit in one 32-group (the layout: element i + cap / 2 is k + 16)
|
||||
if (mm_probe == 1) { // the bench probe: no decode (constant weights)
|
||||
wg = half4(half(n & 7)); wu = wg; wg2 = half4(half(k & 7)); wu2 = wg2;
|
||||
} else if (TK == 32) {
|
||||
// the lane's two runs of this row sit in one 32-group: element i + cap / 2 is (k + 16, n) - part of
|
||||
// the COOP_K4 layout (coop_probe checks it); a runtime check of it here cost 30%
|
||||
W8<F>::run(rgw, k, tab, wg, wg2);
|
||||
W8<F>::run(ruw, k, tab, wu, wu2);
|
||||
} else {
|
||||
|
||||
@@ -10,14 +10,16 @@
|
||||
inline bool coop_k4(mtl::Context& cx) {
|
||||
if (const char* e = std::getenv("COOP_K4")) return std::atoi(e) != 0;
|
||||
const int lib = cx.compile({"mm_bench.metal"}, "#define NO_GGML_COMMON 1\n");
|
||||
mtl::Pipeline p = cx.pipeline(lib, "coop_layout_16_64_32");
|
||||
for (const char* tile : {"coop_layout_16_64_32", "coop_layout_16_32_32"}) {
|
||||
mtl::Pipeline p = cx.pipeline(lib, tile);
|
||||
mtl::Buffer out = cx.alloc(4 * 3 * 32 * 128 * 4);
|
||||
mtl::Dispatch d;
|
||||
d.p = &p;
|
||||
d.args = {mtl::Arg::B(out)};
|
||||
d.tg[0] = 32;
|
||||
cx.run({d});
|
||||
// every lane: elements 4j .. 4j + 3 are (k, n), (k + 1, n), (k + 2, n), (k + 3, n)
|
||||
// every lane: elements 4j .. 4j + 3 are (k, n), (k + 1, n), (k + 2, n), (k + 3, n), and element i + cap / 2 is
|
||||
// (k + 16, n) with k % 32 < 16 (mm_gu3's paired-run decode relies on it)
|
||||
const int* o = out.as<int>();
|
||||
for (int l = 0; l < 32; ++l) {
|
||||
const int* q = o + l * 128 * 4;
|
||||
@@ -25,6 +27,10 @@ inline bool coop_k4(mtl::Context& cx) {
|
||||
for (int i = 0; i + 3 < cap; i += 4)
|
||||
for (int j = 1; j < 4; ++j)
|
||||
if (q[(i + j) * 4] != q[i * 4] + j || q[(i + j) * 4 + 1] != q[i * 4 + 1]) return false;
|
||||
for (int i = 0; i < cap / 2; i += 4)
|
||||
if (q[(i + cap / 2) * 4] != q[i * 4] + 16 || q[(i + cap / 2) * 4 + 1] != q[i * 4 + 1] || (q[i * 4] & 31) >= 16)
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
@@ -33,6 +33,9 @@ int main(int argc, char** argv) {
|
||||
yB = F((size_t)TMAX * HV * S), yh = cx.alloc(2 * (size_t)TMAX * HV * S), qn = F((size_t)TMAX * HK * S),
|
||||
kn = F((size_t)TMAX * HK * S), vc = F((size_t)TMAX * HV * S), gb = F(2 * TMAX * HV), kq = F(TMAX * HK),
|
||||
o = F((size_t)TMAX * HV * S);
|
||||
const size_t NCM = (TMAX + 31) / 32;
|
||||
mtl::Buffer cU = F((size_t)HV * NCM * 32 * S), cW = F((size_t)HV * NCM * 32 * S), clg = F((size_t)HV * NCM * 32),
|
||||
cD = F((size_t)HV * NCM * 32 * S);
|
||||
ref::Rng r(21);
|
||||
for (size_t i = 0; i < (size_t)TMAX * CH; ++i) qkv.as<float>()[i] = r.sym() * 2.0f;
|
||||
for (int i = 0; i < TMAX * HV; ++i) { al.as<float>()[i] = r.sym() * 3.0f; be.as<float>()[i] = r.sym() * 2.0f; }
|
||||
@@ -54,6 +57,22 @@ int main(int argc, char** argv) {
|
||||
d1.grid[1] = T;
|
||||
d1.tg[0] = 256;
|
||||
ds.push_back(d1);
|
||||
if (cb == -100) { // the chunked form: gdn_cprep, gdn_cstate, gdn_cout
|
||||
const uint32_t nc = (T + 31) / 32;
|
||||
static mtl::Pipeline pcp = cx.pipeline(lib, "gdn_cprep"), pcs = cx.pipeline(lib, "gdn_cstate"),
|
||||
pco = cx.pipeline(lib, "gdn_cout");
|
||||
mtl::Dispatch e1; e1.p = &pcp;
|
||||
e1.args = {mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(cU), mtl::Arg::B(cW),
|
||||
mtl::Arg::B(clg), mtl::Arg::V(a)};
|
||||
e1.grid[0] = HV; e1.grid[1] = nc; e1.tg[0] = 256; ds.push_back(e1);
|
||||
mtl::Dispatch e2; e2.p = &pcs;
|
||||
e2.args = {mtl::Arg::B(stB), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(cU), mtl::Arg::B(cW), mtl::Arg::B(clg),
|
||||
mtl::Arg::B(cD), mtl::Arg::B(o), mtl::Arg::V(a)};
|
||||
e2.grid[0] = HV * 4; e2.tg[0] = 256; ds.push_back(e2);
|
||||
mtl::Dispatch e3; e3.p = &pco;
|
||||
e3.args = {mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(cD), mtl::Arg::B(clg), mtl::Arg::B(o), mtl::Arg::V(a)};
|
||||
e3.grid[0] = HV; e3.grid[1] = nc; e3.tg[0] = 128; ds.push_back(e3);
|
||||
} else {
|
||||
mtl::Dispatch d2;
|
||||
d2.p = &scan;
|
||||
d2.args = {mtl::Arg::B(stB), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq),
|
||||
@@ -62,6 +81,7 @@ int main(int argc, char** argv) {
|
||||
d2.grid[0] = cb > 0 ? HV * 4 / cb : HV * 4;
|
||||
d2.tg[0] = 32 * (cb > 0 ? cb : -cb);
|
||||
ds.push_back(d2);
|
||||
}
|
||||
mtl::Dispatch d3;
|
||||
d3.p = &post;
|
||||
d3.args = {mtl::Arg::B(o), mtl::Arg::B(z), mtl::Arg::B(gam), mtl::Arg::B(yh), mtl::Arg::B(yB), mtl::Arg::V(a)};
|
||||
@@ -74,10 +94,10 @@ int main(int argc, char** argv) {
|
||||
int fails = 0;
|
||||
if (!bench)
|
||||
for (int T : {1, 3, 37, 300})
|
||||
for (int cb : {1, 4, -2, -4, -8, -9}) {
|
||||
for (int cb : {1, 4, -2, -4, -8, -9, -100}) {
|
||||
mtl::Constants pc;
|
||||
if (cb < 0 && cb != -9) pc.i(6, 0);
|
||||
mtl::Pipeline scan = cx.pipeline(lib, cb == -9 ? std::string("gdn_scan_v") : cb > 0 ? "gdn_scan_" + std::to_string(cb) : "gdn_scanr_" + std::to_string(-cb), pc);
|
||||
mtl::Pipeline scan = cx.pipeline(lib, cb == -100 ? std::string("gdn_scanr_4") : cb == -9 ? std::string("gdn_scan_v") : cb > 0 ? "gdn_scan_" + std::to_string(cb) : "gdn_scanr_" + std::to_string(-cb), cb == -100 ? mtl::Constants().i(6, 0) : pc);
|
||||
if (cb == -9) cb = -8;
|
||||
mtl::Pipeline step = cx.pipeline(lib, "gdn_step", mtl::Constants().i(2, T));
|
||||
std::copy(st0.as<float>(), st0.as<float>() + S * HV * S, stA.as<float>());
|
||||
@@ -136,14 +156,18 @@ int main(int argc, char** argv) {
|
||||
return median(t);
|
||||
};
|
||||
std::printf("T %d: decode step over the chunk %.2f ms\n", T, timeit({d}) * 1e3);
|
||||
for (const char* nm : {"gdn_scan_v", "gdn_scanr_4", "gdn_scanr_8", "gdn_scan_v", "gdn_scanr_8"}) {
|
||||
const int cb = nm[9] == 'v' ? -8 : nm[8] == 'r' ? -(nm[10] - '0') : nm[9] - '0';
|
||||
for (const char* nm : {"gdn_scanr_4", "chunked", "gdn_scanr_4", "chunked"}) {
|
||||
const bool ch = std::string(nm) == "chunked";
|
||||
const int cb = ch ? -100 : nm[9] == 'v' ? -8 : nm[8] == 'r' ? -(nm[10] - '0') : nm[9] - '0';
|
||||
mtl::Constants pc;
|
||||
if (cb < 0 && nm[9] != 'v') pc.i(6, std::getenv("GDN_PROBE") ? std::atoi(std::getenv("GDN_PROBE")) : 0);
|
||||
mtl::Pipeline scan = cx.pipeline(lib, nm, pc);
|
||||
mtl::Pipeline scan = cx.pipeline(lib, ch ? "gdn_scanr_4" : nm, pc);
|
||||
std::vector<mtl::Dispatch> ds;
|
||||
prompt_ds(T, pre, scan, post, cb, ds);
|
||||
const double all = timeit(ds), tp = timeit({ds[0]}), ts = timeit({ds[1]}), to = timeit({ds[2]});
|
||||
const double all = timeit(ds), tp = timeit({ds[0]}),
|
||||
ts = timeit(std::vector<mtl::Dispatch>(ds.begin() + 1, ds.end() - 1)), to = timeit({ds.back()});
|
||||
if (ch) std::printf(" chunked parts: cprep %.3f, cstate %.3f, cout %.3f ms\n", timeit({ds[1]}) * 1e3, timeit({ds[2]}) * 1e3,
|
||||
timeit({ds[3]}) * 1e3);
|
||||
std::printf("T %d, %-13s: pre %.3f + scan %.3f + post %.3f ms = %.3f ms together\n", T, nm, tp * 1e3, ts * 1e3,
|
||||
to * 1e3, all * 1e3);
|
||||
}
|
||||
|
||||
@@ -93,7 +93,7 @@ int main(int argc, char** argv) {
|
||||
probe.i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0);
|
||||
const bool k4 = coop_k4(cx);
|
||||
if (v3) std::printf("right-input layout: %s\n", k4 ? "runs of 4 k" : "other (per-element decode)");
|
||||
mtl::Pipeline pgu = v3 ? cx.pipeline(lib, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4))
|
||||
mtl::Pipeline pgu = v3 ? cx.pipeline(lib, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0))
|
||||
: cx.pipeline(lib, (v2 ? "mm_gu2_" : "mm_gu_") + std::to_string(gu) + "_" + tl, probe),
|
||||
pdn = v3 ? cx.pipeline(lib, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4))
|
||||
: cx.pipeline(lib, (v2 ? "mm_down2_" : "mm_down_") + std::to_string(dn) + "_" + (v2 ? dtl : tl)),
|
||||
|
||||
@@ -174,6 +174,26 @@ int main(int argc, char** argv) {
|
||||
if (it == pcache.end()) it = pcache.emplace(key, cx.pipeline(l, fn, c)).first;
|
||||
return &it->second;
|
||||
};
|
||||
// ---- SLC prefetch (--prefetch): the next HC read's weights touched concurrently with the first projection GEMV
|
||||
// after the current read, so the latency-bound read finds them in the SLC
|
||||
const bool prefetch = std::getenv("CHAIN_PREFETCH") != nullptr;
|
||||
struct Range { const mtl::Buffer* b; size_t off, bytes; };
|
||||
std::vector<Range> pending;
|
||||
mtl::Buffer sink = cx.alloc(16);
|
||||
const mtl::Pipeline* ptouch = nullptr;
|
||||
auto push_prefetch = [&]() {
|
||||
for (const Range& r : pending) {
|
||||
const uint32_t n16 = (uint32_t)(r.bytes / 16);
|
||||
mtl::Dispatch d;
|
||||
d.p = ptouch;
|
||||
d.args = {mtl::Arg::B(*r.b, r.off), mtl::Arg::B(sink), mtl::Arg::V(n16)};
|
||||
d.grid[0] = (n16 + 1023) / 1024;
|
||||
d.tg[0] = 256;
|
||||
nb = true;
|
||||
push(d, "prefetch", (double)r.bytes);
|
||||
}
|
||||
pending.clear();
|
||||
};
|
||||
// ---- GEMV by tensor, with a tile from the measured table
|
||||
auto gemv2 = [&](const gguf::Tensor& t, const Act& x, const mtl::Buffer& y, size_t y_stride, const std::string& c,
|
||||
int ty_override = -1) {
|
||||
@@ -196,8 +216,10 @@ int main(int argc, char** argv) {
|
||||
d.grid[0] = (n_out + rpt - 1) / rpt;
|
||||
d.tg[0] = S * 32;
|
||||
push(d, c, (double)t.bytes);
|
||||
if (!pending.empty()) push_prefetch();
|
||||
};
|
||||
const mtl::Pipeline* pnorm = Pn(lib, "gr_norm", "gr_norm", {});
|
||||
ptouch = Pn(lib, "slc_touch", "slc_touch", {});
|
||||
const mtl::Pipeline* pdown = Pn(lib, "gr_down_8_8", "gr_down_8_8", mtl::Constants().i(1, ntok));
|
||||
const mtl::Pipeline* pup = Pn(lib, "gr_up", "gr_up", mtl::Constants().i(1, ntok));
|
||||
const mtl::Pipeline* pdownq = Pn(lib, "gr_down_q8_8_8", "gr_down_q8_8_8", mtl::Constants().i(1, ntok));
|
||||
@@ -316,6 +338,28 @@ int main(int argc, char** argv) {
|
||||
d3.tg[0] = 256;
|
||||
push(d3, "hc", (double)tup.bytes);
|
||||
cur_inj = 1 - cur_inj;
|
||||
if (prefetch && !final_) {
|
||||
// the next read: blk.l.hc_attn -> blk.l.hc_ffn -> blk.(l+1).hc_attn ... -> output_hc
|
||||
std::string nx;
|
||||
const size_t at = pre.find(".hc_attn");
|
||||
if (at != std::string::npos) nx = pre.substr(0, at) + ".hc_ffn";
|
||||
else {
|
||||
const int l = std::atoi(pre.c_str() + 4);
|
||||
nx = l == 47 ? std::string("output_hc") : "blk." + std::to_string(l + 1) + ".hc_attn";
|
||||
}
|
||||
const bool fin = nx == "output_hc";
|
||||
for (const std::string& part : fin ? std::vector<std::string>{"_down", "_up"}
|
||||
: std::vector<std::string>{"_down", "_inject", "_up"}) {
|
||||
const gguf::Tensor& tt = T(nx + part + ".weight");
|
||||
if (q83) {
|
||||
uint64_t ne = 1;
|
||||
for (uint64_t dd : tt.shape) ne *= dd;
|
||||
pending.push_back({&Q8, q8off.at(tt.name), (size_t)(ne / 32 * 34) / 16 * 16});
|
||||
} else {
|
||||
pending.push_back({&W, tt.offset, (size_t)tt.bytes / 16 * 16});
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
auto xprep = [&](const mtl::Buffer& x, const Act& out, int n, const std::string& c) {
|
||||
mtl::Dispatch d;
|
||||
|
||||
Reference in new issue
Block a user