diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index 880e863..7a2c9f4 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -1095,6 +1095,11 @@ struct MoeTask { // expert's K dimension per node (banks are already per-row contiguous). constexpr int IBLK = 32; constexpr int HBLK = 32; +// Deduped pass 2 gives one work item the whole H-block for *every* token, so it owns +// those output rows outright and needs no cross-worker reduction. That costs work +// items (n_hblk instead of tokens * n_hblk), so the block is smaller to keep the pool +// fed: H/8 items is ~8 per worker at H=2048 on a 32-core part. +constexpr int HBLK_DD = 8; // -------------------------------- Q4_0 (W4A8) -------------------------------- // Native GGUF Q4_0 experts (gemma4 GGUF): per-32 block = fp16 scale d + 16 packed @@ -1292,7 +1297,23 @@ struct CpuMoeExecutor { std::atomic p2_next{0}; std::atomic prt_next{0}; // ds_fp4 intermediate fp8 round-trip phase int64_t p1_total = 0, p2_total = 0, prt_total = 0; - int n_iblk = 0, n_hblk = 0; + int n_iblk = 0, n_hblk = 0, n_hblk_dd = 0; + // Deduped pass-2 block size. 8 rows keeps the worker pool fed, and for the + // row-major formats the block is just "which output rows", so any size works. + // mxfp4 is different: its bank is transposed, so this becomes mxgemv's `ncol` -- + // the dimension it vectorizes over. Below 16 the whole call drops into mxgemv's + // scalar tail, which measured 2.7x SLOWER than not deduping at all. Keep the full + // block there. + int hblk_dd = HBLK_DD; + // Expert dedup for pass 1 (see build_dedup). dd_route holds the (tok*top_k + k) + // route ids grouped by expert; dd_expert[j] owns dd_route[dd_start[j] .. + // dd_start[j+1]). n_dd == 0 means "not deduped, use the per-route work split". + std::vector dd_expert, dd_start, dd_route, dd_hist; + int n_dd = 0; + // Read once per executor rather than once per process: a static would freeze the + // first value seen, which silently disables the A/B when both settings are + // exercised from one process. + bool dedup_enabled = true; std::atomic done_count{0}; std::atomic bar_count{0}; std::atomic bar_sense{0}; @@ -1386,6 +1407,7 @@ struct CpuMoeExecutor { // W4A8 (activations pre-quantized to Q8_0); select_q4dot picks VPDPBUSD / VPMADDUBSW // / scalar for the tier, so the tag reflects which of those q4dot resolved to. nvi8dot = select_nvi8dot(); + if (const char* de = getenv("FREETOKEN_CPU_MOE_DEDUP")) dedup_enabled = de[0] != '0'; use_vnni = (weight_format == WF_NVFP4) && (nvi8dot != nullptr); use_q4a8 = (weight_format == WF_Q4_0); const char* q4tag = use_q4a8 ? (cpu_has_avxvnni() ? "+vnni(q4_0-w4a8)" : "+q4_0-w4a8") : ""; @@ -1571,13 +1593,109 @@ struct CpuMoeExecutor { } } + // Group this task's routes by expert, counting-sort style. Only worth it when more + // than one token can collide on an expert, and only for the formats that go through + // gemm1_dot (mxfp4 and ds_fp4 own their pass-1 bodies). Sets n_dd = 0 to opt out. + void build_dedup(const MoeTask* t) { + n_dd = 0; + if (!dedup_enabled || t->num_tokens < 2) return; + const int routes = t->num_tokens * top_k; + dd_hist.assign(num_experts, 0); + int valid = 0; + for (int r = 0; r < routes; ++r) { + const int e = t->ids[r]; + if (e < 0 || e >= num_experts) continue; + ++dd_hist[e]; + ++valid; + } + if (valid == 0) return; + dd_expert.clear(); + dd_start.clear(); + dd_expert.reserve(valid); + dd_start.reserve(valid + 1); + int run = 0; + for (int e = 0; e < num_experts; ++e) { + if (!dd_hist[e]) continue; + dd_expert.push_back(e); + dd_start.push_back(run); + run += dd_hist[e]; + dd_hist[e] = dd_start.back(); // reuse as the per-expert write cursor + } + dd_start.push_back(run); + dd_route.resize(valid); + for (int r = 0; r < routes; ++r) { + const int e = t->ids[r]; + if (e < 0 || e >= num_experts) continue; + dd_route[dd_hist[e]++] = r; + } + // Every expert distinct -> the split is the same work, so keep the simpler path. + if ((int)dd_expert.size() == valid) return; + n_dd = (int)dd_expert.size(); + if (getenv("FREETOKEN_CPU_MOE_DEDUP_DEBUG")) + fprintf(stderr, "[dedup] tokens=%d routes=%d valid=%d unique=%d reuse=%.2fx\n", + t->num_tokens, routes, valid, n_dd, (double)valid / n_dd); + } + + // Pass 1 over (unique expert, row block): the expert's gate and up rows are loaded + // once and reused across every token routed to it. Both rows stay in L1 across the + // inner route loop (2 * H * 2 bytes = 16 KiB at H=4096), so the repeat reads never + // reach DRAM -- which is the whole point, the GEMV being DRAM-bound. + void do_pass1_dedup(const MoeTask* t, int64_t p) { + const int64_t ib = p % n_iblk; + const int es = static_cast(p / n_iblk); + const int e = dd_expert[es]; + const int r0 = dd_start[es], r1 = dd_start[es + 1]; + const bf16_t* gate_up_l = reinterpret_cast(tbl_at(gate_up_tbl, t->layer_id)); + const uint8_t* gu_packed_l = reinterpret_cast(tbl_at(gate_up_tbl, t->layer_id)); + const uint8_t* gu_scale_l = reinterpret_cast(tbl_at(gu_scale_tbl, t->layer_id)); + const uint16_t* gu_global_l = + reinterpret_cast(tbl_at(gu_global_tbl, t->layer_id)); + const int i0 = static_cast(ib) * IBLK; + const int i1 = std::min(I, i0 + IBLK); + const bool swigluoai = act == ACT_SWIGLUOAI; + const float lim = swiglu_limit, alpha = swiglu_alpha; + for (int i = i0; i < i1; ++i) { + for (int r = r0; r < r1; ++r) { + const int route = dd_route[r]; + const int tok = route / top_k; + const float w_in = apply_on_input ? t->w[route] : 1.0f; + const bf16_t* x_row = t->x + (size_t)tok * H; + const float* xe = needs_di ? xe_scratch.data() + (size_t)tok * (H / 2) : nullptr; + const float* xo = needs_di ? xo_scratch.data() + (size_t)tok * (H / 2) : nullptr; + const int8_t* xi8 = + (use_vnni || use_q4a8) ? xi8_scratch.data() + (size_t)tok * H : nullptr; + const float* xas = use_vnni ? xas_scratch.data() + (size_t)tok * (H / 16) + : use_q4a8 ? xas_scratch.data() + (size_t)tok * (H / 32) + : nullptr; + float gate = gemm1_dot(gate_up_l, gu_packed_l, gu_scale_l, gu_global_l, e, i, + x_row, xe, xo, xi8, xas) * w_in; + float up = gemm1_dot(gate_up_l, gu_packed_l, gu_scale_l, gu_global_l, e, I + i, + x_row, xe, xo, xi8, xas) * w_in; + bf16_t* g_row = g_scratch.data() + (size_t)route * I; + if (swigluoai) { + if (gate > lim) gate = lim; + if (up > lim) up = lim; + else if (up < -lim) up = -lim; + const float glu = gate / (1.0f + std::exp(-gate * alpha)); + g_row[i] = f32_to_bf16(glu * (up + 1.0f)); + } else { + g_row[i] = f32_to_bf16(act_apply(act, gate) * up); + } + } + } + } + void do_pass1(const MoeTask* t, int64_t p) { if (fmt == WF_MXFP4) { - do_pass1_mxfp4(t, p); + if (n_dd > 0) do_pass1_mxfp4_dedup(t, p); else do_pass1_mxfp4(t, p); return; } if (fmt == WF_DSFP4) { - do_pass1_dsfp4(t, p); + if (n_dd > 0) do_pass1_dsfp4_dedup(t, p); else do_pass1_dsfp4(t, p); + return; + } + if (n_dd > 0) { + do_pass1_dedup(t, p); return; } const int64_t ib = p % n_iblk; @@ -1626,13 +1744,69 @@ struct CpuMoeExecutor { } } + // Pass 2 over an H-block, all tokens, all experts. Deduping pass 2 the way pass 1 + // is done -- work item per (expert, block) -- would have several experts summing + // into the same y row and need a reduction. Giving one work item every token for + // its rows sidesteps that: the rows are exclusively owned, the accumulation happens + // in a private fp32 buffer, and each expert's down rows are still read once and + // reused across the tokens routed to it. + // + // Summation order changes (expert order rather than route order), so the fp32 + // rounding differs in the last bits from the non-deduped path -- the same latitude + // the kernel already takes between its ISA tiers. + void do_pass2_dedup(const MoeTask* t, int64_t p) { + const int h0 = static_cast(p) * hblk_dd; + const int h1 = std::min(H, h0 + hblk_dd); + if (h0 >= h1) return; + const int nh = h1 - h0, nt = t->num_tokens; + thread_local std::vector acc; + acc.assign((size_t)nt * nh, 0.0f); + + const bf16_t* down_l = reinterpret_cast(tbl_at(down_tbl, t->layer_id)); + const uint8_t* dn_packed_l = reinterpret_cast(tbl_at(down_tbl, t->layer_id)); + const uint8_t* dn_scale_l = reinterpret_cast(tbl_at(dn_scale_tbl, t->layer_id)); + const uint16_t* dn_global_l = + reinterpret_cast(tbl_at(dn_global_tbl, t->layer_id)); + + for (int es = 0; es < n_dd; ++es) { + const int e = dd_expert[es]; + const int r0 = dd_start[es], r1 = dd_start[es + 1]; + for (int h = h0; h < h1; ++h) { + for (int r = r0; r < r1; ++r) { + const int route = dd_route[r]; + const int tok = route / top_k; + const float w_out = apply_on_input ? 1.0f : t->w[route]; + const size_t gr = (size_t)route; + const bf16_t* g_row = g_scratch.data() + gr * I; + const float* ge = needs_di ? ge_scratch.data() + gr * (I / 2) : nullptr; + const float* go = needs_di ? go_scratch.data() + gr * (I / 2) : nullptr; + const int8_t* gi8 = (use_vnni || use_q4a8) ? gi8_scratch.data() + gr * I : nullptr; + const float* gas = use_vnni ? gas_scratch.data() + gr * (I / 16) + : use_q4a8 ? gas_scratch.data() + gr * (I / 32) + : nullptr; + acc[(size_t)tok * nh + (h - h0)] += + gemm2_dot(down_l, dn_packed_l, dn_scale_l, dn_global_l, e, h, g_row, ge, go, + gi8, gas) * w_out; + } + } + } + for (int tok = 0; tok < nt; ++tok) { + bf16_t* y_row = t->y + (size_t)tok * H; + for (int h = h0; h < h1; ++h) y_row[h] = f32_to_bf16(acc[(size_t)tok * nh + (h - h0)]); + } + } + void do_pass2(const MoeTask* t, int64_t p) { if (fmt == WF_MXFP4) { - do_pass2_mxfp4(t, p); + if (n_dd > 0) do_pass2_fp4_dedup(t, p, /*mx=*/true); else do_pass2_mxfp4(t, p); return; } if (fmt == WF_DSFP4) { - do_pass2_dsfp4(t, p); + if (n_dd > 0) do_pass2_fp4_dedup(t, p, /*mx=*/false); else do_pass2_dsfp4(t, p); + return; + } + if (n_dd > 0) { + do_pass2_dedup(t, p); return; } const int64_t hb = p % n_hblk; @@ -1676,6 +1850,131 @@ struct CpuMoeExecutor { // Dequant: w = E2M1[code] * 2^(e8m0_scale - 127); two codes per byte (low nibble // first), one e8m0 scale per 32 contiguous K. Matches kernel/triton/mxfp4_moe.py. + // mxfp4 pass 1, deduped. mxgemv computes a whole (Hh x ncol) tile for one token, so + // reuse here is the tile staying resident across the expert's routes -- 128 KiB at + // H=4096, so L2 rather than L1, but still not DRAM. + void do_pass1_mxfp4_dedup(const MoeTask* t, int64_t p) { + const int64_t ib = p % n_iblk; + const int es = static_cast(p / n_iblk); + const int e = dd_expert[es]; + const int r0 = dd_start[es], r1 = dd_start[es + 1]; + const uint8_t* gu_packed_l = reinterpret_cast(tbl_at(gate_up_tbl, t->layer_id)); + const uint8_t* gu_scale_l = reinterpret_cast(tbl_at(gu_scale_tbl, t->layer_id)); + const bf16_t* gu_bias_l = reinterpret_cast(tbl_at(gu_bias_tbl, t->layer_id)); + const int N2 = 2 * I, Hh = H / 2; + const int i0 = static_cast(ib) * IBLK; + const int i1 = std::min(I, i0 + IBLK); + const int nunit = i1 - i0, col0 = 2 * i0, ncol = 2 * nunit; + const uint8_t* blk_e = gu_packed_l + (size_t)e * Hh * N2; + const uint8_t* scl_e = gu_scale_l + (size_t)e * (size_t)(H / 32) * N2; + const bf16_t* bias_e = gu_bias_l + (size_t)e * N2 + col0; + const float lim = swiglu_limit, alpha = swiglu_alpha; + for (int r = r0; r < r1; ++r) { + const int route = dd_route[r]; + const int tok = route / top_k; + float gu[2 * IBLK]; + mxgemv(gu, blk_e + col0, scl_e + col0, t->x + (size_t)tok * H, Hh, N2, ncol, + e2m1_lut, e8m0_lut); + bf16_t* g_row = g_scratch.data() + (size_t)route * I; + for (int j = 0; j < nunit; ++j) { + float gate = gu[2 * j] + bf16_to_f32(bias_e[2 * j]); + float up = gu[2 * j + 1] + bf16_to_f32(bias_e[2 * j + 1]); + if (gate > lim) gate = lim; + if (up > lim) up = lim; + else if (up < -lim) up = -lim; + const float glu = gate / (1.0f + std::exp(-gate * alpha)); + g_row[i0 + j] = f32_to_bf16(glu * (up + 1.0f)); + } + } + } + + // ds_fp4 pass 1, deduped: the expert's gate and up rows are read once and reused + // across its routes, exactly as in the generic path. + void do_pass1_dsfp4_dedup(const MoeTask* t, int64_t p) { + const int64_t ib = p % n_iblk; + const int es = static_cast(p / n_iblk); + const int e = dd_expert[es]; + const int r0 = dd_start[es], r1 = dd_start[es + 1]; + const uint8_t* gu_packed_l = reinterpret_cast(tbl_at(gate_up_tbl, t->layer_id)); + const uint8_t* gu_scale_l = reinterpret_cast(tbl_at(gu_scale_tbl, t->layer_id)); + const int N2 = 2 * I, Hh = H / 2, Hs = H / 32; + const uint8_t* gp = gu_packed_l + (size_t)e * N2 * Hh; + const uint8_t* gs = gu_scale_l + (size_t)e * N2 * Hs; + const int i0 = static_cast(ib) * IBLK; + const int i1 = std::min(I, i0 + IBLK); + const float lim = swiglu_limit; + for (int i = i0; i < i1; ++i) { + for (int r = r0; r < r1; ++r) { + const int route = dd_route[r]; + const int tok = route / top_k; + const float* xe = xe_scratch.data() + (size_t)tok * (H / 2); + const float* xo = xo_scratch.data() + (size_t)tok * (H / 2); + float gate = bf16_to_f32(f32_to_bf16( + dsdot(gp + (size_t)i * Hh, gs + (size_t)i * Hs, xe, xo, H, e2m1_lut, e8m0_lut))); + float up = bf16_to_f32(f32_to_bf16(dsdot( + gp + (size_t)(I + i) * Hh, gs + (size_t)(I + i) * Hs, xe, xo, H, e2m1_lut, e8m0_lut))); + if (lim > 0.0f) { + if (gate > lim) gate = lim; + if (up > lim) up = lim; + else if (up < -lim) up = -lim; + } + const float glu = gate / (1.0f + std::exp(-gate)); + g_scratch[(size_t)route * I + i] = f32_to_bf16(glu * up); + } + } + } + + // Shared deduped pass 2 for both fp4 formats: one work item owns an H-block for + // every token (see do_pass2_dedup for why), accumulating in a private fp32 buffer. + void do_pass2_fp4_dedup(const MoeTask* t, int64_t p, bool mx) { + const int h0 = static_cast(p) * hblk_dd; + const int h1 = std::min(H, h0 + hblk_dd); + if (h0 >= h1) return; + const int nh = h1 - h0, nt = t->num_tokens; + thread_local std::vector acc; + acc.assign((size_t)nt * nh, 0.0f); + const uint8_t* dn_packed_l = reinterpret_cast(tbl_at(down_tbl, t->layer_id)); + const uint8_t* dn_scale_l = reinterpret_cast(tbl_at(dn_scale_tbl, t->layer_id)); + const bf16_t* dn_bias_l = + mx ? reinterpret_cast(tbl_at(dn_bias_tbl, t->layer_id)) : nullptr; + const int Ih = I / 2, Is = I / 32; + for (int es = 0; es < n_dd; ++es) { + const int e = dd_expert[es]; + const int r0 = dd_start[es], r1 = dd_start[es + 1]; + for (int r = r0; r < r1; ++r) { + const int route = dd_route[r]; + const int tok = route / top_k; + const float wt = t->w[route]; + if (mx) { + const uint8_t* blk_e = dn_packed_l + (size_t)e * Ih * H; + const uint8_t* scl_e = dn_scale_l + (size_t)e * (size_t)Is * H; + float part[HBLK]; + mxgemv(part, blk_e + h0, scl_e + h0, g_scratch.data() + (size_t)route * I, Ih, H, + nh, e2m1_lut, e8m0_lut); + const bf16_t* bias_e = dn_bias_l + (size_t)e * H + h0; + for (int c = 0; c < nh; ++c) + acc[(size_t)tok * nh + c] += (part[c] + bf16_to_f32(bias_e[c])) * wt; + } else { + const uint8_t* dp_e = dn_packed_l + (size_t)e * (size_t)H * Ih; + const uint8_t* ds_e = dn_scale_l + (size_t)e * (size_t)H * Is; + const float* ge = ge_scratch.data() + (size_t)route * (I / 2); + const float* go = go_scratch.data() + (size_t)route * (I / 2); + for (int c = 0; c < nh; ++c) { + const int h = h0 + c; + // the reference rounds each route's weighted output to bf16 before summing + acc[(size_t)tok * nh + c] += bf16_to_f32(f32_to_bf16( + dsdot(dp_e + (size_t)h * Ih, ds_e + (size_t)h * Is, ge, go, I, e2m1_lut, + e8m0_lut) * wt)); + } + } + } + } + for (int tok = 0; tok < nt; ++tok) { + bf16_t* y_row = t->y + (size_t)tok * H; + for (int c = 0; c < nh; ++c) y_row[h0 + c] = f32_to_bf16(acc[(size_t)tok * nh + c]); + } + } + void do_pass1_mxfp4(const MoeTask* t, int64_t p) { const int64_t ib = p % n_iblk; const int64_t tk = p / n_iblk; @@ -1880,6 +2179,7 @@ struct CpuMoeExecutor { } void submit(MoeTask* t) { + build_dedup(t); n_iblk = (I + IBLK - 1) / IBLK; n_hblk = (H + HBLK - 1) / HBLK; // Grow the per-token intermediate scratch if a larger batch shows up than the @@ -1887,8 +2187,15 @@ struct CpuMoeExecutor { // this happens at most once, before any capture, while the pool is idle). const size_t need = static_cast(t->num_tokens) * top_k * I; if (need > g_scratch.size()) g_scratch.resize(need); - p1_total = static_cast(t->num_tokens) * top_k * n_iblk; - p2_total = static_cast(t->num_tokens) * n_hblk; + // Deduped: one work item per (unique expert, output-row block) instead of per + // (token, route, block), so an expert's rows are read from DRAM once for every + // token routed to it rather than once per token. + p1_total = n_dd > 0 ? static_cast(n_dd) * n_iblk + : static_cast(t->num_tokens) * top_k * n_iblk; + hblk_dd = (fmt == WF_MXFP4) ? HBLK : HBLK_DD; + n_hblk_dd = (H + hblk_dd - 1) / hblk_dd; + p2_total = n_dd > 0 ? static_cast(n_hblk_dd) + : static_cast(t->num_tokens) * n_hblk; prt_total = (needs_di || use_q4a8) ? static_cast(t->num_tokens) * top_k : 0; p1_next.store(0, std::memory_order_relaxed); p2_next.store(0, std::memory_order_relaxed);