From 57e4d4dae88f6af85f81c230796ee9f2db278801 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20Edson?= Date: Tue, 27 Oct 2026 20:37:04 -0300 Subject: [PATCH] perf: metal: moe_gu decodes gate rows then up rows for 3- and 4-token groups (MOE_SEQ) Function constant 25 (MOE_SEQ) adds gu_body_seq: for the 3- and 4-entry groups of a speculative window, an expert's gate units are decoded and accumulated first, then its up units, in one pass over the weights with the same math order. The host sets it only for windows of 3+ tokens, gr 1 and Q2_0, IQ2_XXS or IQ3_XXS gate/up; 1- and 2-token pipelines are built without the path. expert_parity takes MOE_SEQ and PARITY_SHARE from the environment. --- metal/kernels/moe.metal | 59 +++++++++++++++++++++++++++++++++-- metal/src/engine/window.cpp | 8 ++++- metal/tests/expert_parity.cpp | 4 +-- 3 files changed, 65 insertions(+), 6 deletions(-) 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());