diff --git a/metal/kernels/moe.metal b/metal/kernels/moe.metal index e3f0b1d68ff..61eadfa2336 100644 --- a/metal/kernels/moe.metal +++ b/metal/kernels/moe.metal @@ -199,9 +199,54 @@ static inline void gu_body(device const uchar* g0, device const uchar* u0, devic } } +// NE >= 3: the gate rows' units, then the up rows' units, so only one of them is decoded at a time (M1 Max: the +// joint body spills at 3-4 entries, 2.3x the 2-entry time for 15% more weight bytes) +template +static inline void gu_body_seq(device const uchar* g0, device const uchar* u0, device const float* x, + device const float* xs16, device float* h, device const int* etok, int c, int row0, int li, + constant MoeArgs& a, threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT; + float ag[NR][NE], au[NR][NE]; + for (int r = 0; r < NR; ++r) + for (int t = 0; t < NE; ++t) ag[r][t] = au[r][t] = 0.0f; + const int U = int(a.n_embd) / UNIT; + for (int u = li; u < U; u += LPR) { + { + float q[NR][UNIT]; + float4 cq[NR]; + for (int r = 0; r < NR; ++r) Dec::run(g0 + r * a.gu_row, u, q[r], cq[r], tab); + for (int t = 0; t < NE; ++t) { + float xs[UNIT], xsum[2]; + X::load(x + etok[t] * int(a.n_embd), xs16 + etok[t] * int(a.n_embd / 16), u, xs, xsum); + for (int r = 0; r < NR; ++r) ag[r][t] += unit_combine(q[r], cq[r], xs, xsum); + } + } + { + float q[NR][UNIT]; + float4 cq[NR]; + for (int r = 0; r < NR; ++r) Dec::run(u0 + r * a.gu_row, u, q[r], cq[r], tab); + for (int t = 0; t < NE; ++t) { + float xs[UNIT], xsum[2]; + X::load(x + etok[t] * int(a.n_embd), xs16 + etok[t] * int(a.n_embd / 16), u, xs, xsum); + for (int r = 0; r < NR; ++r) au[r][t] += unit_combine(q[r], cq[r], xs, xsum); + } + } + } + for (int r = 0; r < NR; ++r) + for (int t = 0; t < NE; ++t) { + const float gv = lane_sum(ag[r][t]), uv = lane_sum(au[r][t]); + if (li == 0) h[(c + t) * int(a.n_ff) + row0 + r] = gv / (1.0f + exp(-gv)) * uv; + } +} + // bench probe (function constant 22): 1 = no table load (wrong values for the i-quants) constant int MOE_PROBE [[function_constant(22)]]; constant int moe_probe = is_function_constant_defined(MOE_PROBE) ? MOE_PROBE : 0; +// MOE_SEQ (function constant 25): the 3- and 4-entry groups take gu_body_seq. The host sets it only for windows of 3 or +// more tokens, so the pipelines of 1- and 2-token windows compile without that path (with it in the same pipeline the +// M1 Max ran 1 and 2 tokens 8-10% slower) +constant bool MOE_SEQ [[function_constant(25)]]; +constant bool moe_seq = is_function_constant_defined(MOE_SEQ) && MOE_SEQ; template kernel void moe_gu(device const uchar* W [[buffer(0)]], device const float* x [[buffer(1)]], @@ -230,12 +275,20 @@ kernel void moe_gu(device const uchar* W [[buffer(0)]], device const float* x [[ switch (ne) { case 1: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; case 2: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; - case 3: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; - case 4: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; + case 3: + if (moe_seq) gu_body_seq(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); + else gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); + break; + case 4: + if (moe_seq) gu_body_seq(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); + else gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); + break; default: for (int k = 0; k < ne; k += 4) { const int m = min(4, ne - k); - if (m == 4) gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + if (m == 4 && moe_seq) gu_body_seq(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + else if (m == 4) gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + else if (m == 3 && moe_seq) gu_body_seq(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); else if (m == 3) gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); else if (m == 2) gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); else gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); diff --git a/metal/src/engine/window.cpp b/metal/src/engine/window.cpp index add35c0e9f3..a72543f8bc1 100644 --- a/metal/src/engine/window.cpp +++ b/metal/src/engine/window.cpp @@ -974,7 +974,13 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 mtl::Constants ac = stream_ ? mtl::Constants().b(20, true) : mtl::Constants(); ac.b(23, ct); const std::string ak = std::string(stream_ ? "#arena" : "") + (ct ? "#ct" : ""); - g.p = P(lib_, gun, ac, gun + ak); + // function constant 25 (MOE_SEQ): windows of 3+ tokens decode an expert's gate rows, then its up rows; the 1- and + // 2-token pipelines are built without that path (in the same pipeline it slowed them 8-10% on the M1 Max). Only + // Q2_0, IQ2_XXS and IQ3_XXS gate/up with 1 row a lane: on the M1 Max the K-quants, Q8_0, IQ4_XS, IQ2_XS and the 2-row Q2_0 tile ran 10-43% slower with it + const bool seq = ntok >= 3 && gr == 1 && (tg_.type == ref::Q2_0 || tg_.type == ref::IQ2_XXS || tg_.type == ref::IQ3_XXS); + mtl::Constants gc = ac; + if (seq) gc.b(25, true); + g.p = P(lib_, gun, gc, gun + ak + (seq ? "#seq" : "")); const bool hg = had_gu(l), hd = had_dn(l); if (hg) had_rows(ds, lib_, "d:", XF_.b, 0, XR_, 0, ntok, N, hsx_.handle ? &hsx_ : nullptr, false); // before the host's turn g.args = {mtl::Arg::B(EW, ebase), hg ? mtl::Arg::B(XR_) : mtl::Arg::B(XF_.b, XF_.plane_bytes() * xkind(tg_.type)), mtl::Arg::B(h_), diff --git a/metal/tests/expert_parity.cpp b/metal/tests/expert_parity.cpp index 5a596a74ee8..946e94fb04a 100644 --- a/metal/tests/expert_parity.cpp +++ b/metal/tests/expert_parity.cpp @@ -414,7 +414,7 @@ int main(int argc, char** argv) { K.plan = cx.pipeline(lib, "moe_plan"); K.prep = cx.pipeline(lib, "xprep"); K.comb = cx.pipeline(lib, "moe_combine"); - K.gu = cx.pipeline(lib, kname("moe_gu", gu, a), mtl::Constants().i(22, std::getenv("MOE_PROBE") ? std::atoi(std::getenv("MOE_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr)); + K.gu = cx.pipeline(lib, kname("moe_gu", gu, a), mtl::Constants().i(22, std::getenv("MOE_PROBE") ? std::atoi(std::getenv("MOE_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr).b(25, std::getenv("MOE_SEQ") != nullptr)); K.dn = cx.pipeline(lib, kname("moe_down", dn == ref::Q2_0 && !g_dn_prep ? 142 : dn, b)); K.cgu = a; K.cdn = b; @@ -491,7 +491,7 @@ int main(int argc, char** argv) { } std::vector ids; std::vector w; - random_routing(rng, ntok, NE, ids, w, 0.3); + random_routing(rng, ntok, NE, ids, w, std::getenv("PARITY_SHARE") ? std::atof(std::getenv("PARITY_SHARE")) : 0.3); mtl::Buffer bid = cx.alloc(4 * ids.size()), bw = cx.alloc(4 * w.size()), plan = cx.alloc(4 * P_INTS); std::memcpy(bid.ptr, ids.data(), 4 * ids.size()); std::memcpy(bw.ptr, w.data(), 4 * w.size());