From 49d8a0f1b523fa31c81118a67cdbf55693357813 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 22 Aug 2026 13:34:11 -0400 Subject: [PATCH 1/3] perf(cpu-moe): dedup experts across the batch in pass 1 Pass 1 split its work by (token, route, row block), so when two tokens in a decode batch routed to the same expert, that expert's gate_up rows were streamed from DRAM twice. The GEMV is DRAM-bound, so that is the whole cost. Group the routes by expert (counting sort over the task's ids) and split pass 1 by (unique expert, row block) instead, with the routes for that expert as the inner loop. Each weight row is then read from memory once and reused from L1 across every token routed to it -- both the gate row and the up row are 4 KiB at H=4096, so 8 KiB stays resident across the inner loop. No kernel changes: the reuse comes from loop order, so every format on the gemm1_dot path (bf16, nvfp4, fp8_block, q4_0) benefits at once. Measured on 2x Xeon Gold 6526Y, E=256, top_k=8, H=2048, I=768, uniform-random routing (the pessimistic case -- real routing is skewed, so collisions are more common): bs routes unique reuse off on delta 4 32 30 1.07x 3.54ms 3.39ms +4.3% 8 64 56 1.14x 7.12ms 6.80ms +4.6% 16 128 104 1.23x 14.51ms 13.07ms +11.0% 32 256 156 1.64x 29.97ms 21.97ms +36.4% 64 512 213 2.40x 59.12ms 35.44ms +66.8% Those track a simple traffic model to within 2%: pass 1 is about two thirds of the bytes (gate_up is [2I, H] against down's [H, I]), so the expected speedup is 1 / (2/3 / reuse + 1/3) -- 1.35x at bs=32 and 1.64x at bs=64 against 1.36x and 1.67x measured. Inert below bs=2 and skipped when every route already has a distinct expert, so single-stream decode keeps exactly the old work split. mxfp4 and ds_fp4 own their pass-1 bodies and are untouched; they would follow the same shape. Pass 2 is deliberately not deduped. Its work items are per-token and own their output rows exclusively; deduping it would have several experts accumulating into the same y row and needs a reduction, which is a separate change. `FREETOKEN_CPU_MOE_DEDUP=0` restores the old split, and `FREETOKEN_CPU_MOE_DEDUP_DEBUG=1` reports the reuse factor -- worth having, since a first attempt at the toggle cached the env in a function-local static and silently disabled both arms of the A/B, which read as "dedup does nothing" (+1.4%). Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01GpGe2fQ5pDShGrnSuksfun --- .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 113 +++++++++++++++++- 1 file changed, 112 insertions(+), 1 deletion(-) 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..72c5c7a 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -1293,6 +1293,15 @@ struct CpuMoeExecutor { 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; + // 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 +1395,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,7 +1581,103 @@ 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 || fmt == WF_MXFP4 || fmt == WF_DSFP4) 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 (n_dd > 0) { + do_pass1_dedup(t, p); + return; + } if (fmt == WF_MXFP4) { do_pass1_mxfp4(t, p); return; @@ -1880,6 +1986,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,7 +1994,11 @@ 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; + // 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; p2_total = 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); From a60a654cf3e7c97d434152e9459918688559a4aa Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 22 Aug 2026 13:39:39 -0400 Subject: [PATCH 2/3] perf(cpu-moe): dedup experts in pass 2 as well Pass 1 dedup left the down projection reading each expert once per route. Deduping it the same way -- a work item per (expert, row block) -- would have several experts summing into the same y row and need a cross-worker reduction. Give one work item every token for its H-block instead. The rows are then owned outright, the accumulation happens in a private fp32 buffer, and an expert's down rows are still read once and reused across the tokens routed to it. The block is HBLK_DD = 8 rather than 32 because the item count drops from tokens * n_hblk to n_hblk, and H/8 keeps ~8 items per worker at H=2048 on a 32-core part. On top of pass 1 (same rig: 2x Xeon Gold 6526Y, E=256, top_k=8, H=2048, I=768, uniform-random routing): bs reuse off pass1 pass1+pass2 8 1.14x 7.27ms +4.6% +11.8% 16 1.23x 14.63ms +11.0% +25.3% 32 1.64x 29.94ms +36.4% +78.8% 64 2.40x 59.18ms +66.8% +158.2% 2.58x at bs=64, slightly ahead of the 2.40x the traffic model predicts -- the smaller H-block also helps locality. Numerics: the fp32 accumulation now runs in expert order rather than route order, so the last bits differ from the non-deduped path. That is the same latitude the kernel already takes between ISA tiers, and the GPU-comparison test (bs 1/2/5/16) covers it. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01GpGe2fQ5pDShGrnSuksfun --- .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 67 ++++++++++++++++++- 1 file changed, 65 insertions(+), 2 deletions(-) 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 72c5c7a..55296c8 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,7 @@ 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; // 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". @@ -1732,7 +1737,63 @@ 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 (n_dd > 0) { + do_pass2_dedup(t, p); + return; + } if (fmt == WF_MXFP4) { do_pass2_mxfp4(t, p); return; @@ -1999,7 +2060,9 @@ struct CpuMoeExecutor { // 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; - p2_total = static_cast(t->num_tokens) * n_hblk; + 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); From 4a6d86bf5e30630d9f0da441fb5c9c8ebe48020e Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 22 Aug 2026 13:47:01 -0400 Subject: [PATCH 3/3] perf(cpu-moe): extend expert dedup to ds_fp4 and mxfp4 Both fp4 formats own their pass bodies, so neither got anything from the generic dedup. Give them the same treatment: pass 1 splits by (unique expert, row block) with the expert's routes inner, and pass 2 shares one deduped body -- one work item owns an H-block for every token, accumulating in a private fp32 buffer. ds_fp4 reuses its gate/up rows from L1 exactly as the generic path does. mxfp4's mxgemv computes a whole tile per token, so its reuse is the tile staying resident across the expert's routes -- 128 KiB at H=4096, L2 rather than L1, still not DRAM. ds_fp4 (E=128 H=4096 I=2048 top_k=6) mxfp4 (E=64 H=2880 I=2880 top_k=4) bs reuse off on delta bs reuse off on delta 8 1.20x 9.24ms 7.21ms +28.2% 8 1.39x 7.43ms 5.57ms +33.3% 16 1.35x 17.44ms 12.89ms +35.2% 16 1.45x 13.68ms 10.72ms +27.6% 32 1.94x 33.86ms 19.62ms +72.5% 32 2.33x 26.57ms 14.70ms +80.7% The deduped pass-2 block size has to be format-aware, and getting it wrong is expensive. For the row-major formats the block is just "which output rows", so 8 rows is free and keeps the worker pool fed now that the item count has dropped from tokens * n_hblk to n_hblk. mxfp4's bank is transposed, so the same number becomes mxgemv's `ncol` -- the dimension it vectorizes over -- and below 16 every call falls into mxgemv's scalar tail. Measured at HBLK_DD = 8, mxfp4 dedup ran 2.7x SLOWER than not deduping (-63% at bs=8, -62% at bs=32). It keeps the full 32-row block. Covered by the existing GPU-comparison tests, which run these formats at batch sizes where dedup engages: test_cpu_decode_mxfp4_matches_gpu_splitk at bs 1/4/8 and test_cpu_decode_dsfp4_matches_gpu at bs 1/3/8. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01GpGe2fQ5pDShGrnSuksfun --- .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 165 ++++++++++++++++-- 1 file changed, 149 insertions(+), 16 deletions(-) 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 55296c8..7a2c9f4 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -1298,6 +1298,13 @@ struct CpuMoeExecutor { 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, 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". @@ -1591,7 +1598,7 @@ struct CpuMoeExecutor { // 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 || fmt == WF_MXFP4 || fmt == WF_DSFP4) return; + 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; @@ -1679,16 +1686,16 @@ struct CpuMoeExecutor { } void do_pass1(const MoeTask* t, int64_t p) { - if (n_dd > 0) { - do_pass1_dedup(t, p); - return; - } 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; @@ -1748,8 +1755,8 @@ struct CpuMoeExecutor { // 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); + 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; @@ -1790,16 +1797,16 @@ struct CpuMoeExecutor { } void do_pass2(const MoeTask* t, int64_t p) { - if (n_dd > 0) { - do_pass2_dedup(t, p); - return; - } 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; @@ -1843,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; @@ -2060,7 +2192,8 @@ struct CpuMoeExecutor { // 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; - n_hblk_dd = (H + HBLK_DD - 1) / HBLK_DD; + 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;