mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-07 18:35:04 -10:00
metal: chunked GDN scan (WY form, chunks of 32) - correct, slower in scalar fp32 than gdn_scan_r; kept, not used
- gdn_prompt.metal: gdn_cprep (per head and chunk: gates in log space, A, forward substitution for U and W), gdn_cstate (per head and 32 state columns: the chunks in order), gdn_cout (the intra-chunk output) - gdn_prompt_parity: CB -100 = the chunked form; y 3e-7, state 2e-7 vs the decode step; 7.4 ms vs 1.9 ms at T 2048 (cprep 1.5, cstate 4.7, cout 1.4) - it needs the matrix units to pay off
This commit is contained in:
Niko1221 committed
1 parent
072c19f60f
commit
272bbbbd4a
2 files changed
+192
-6
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;
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user