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/4] 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 880e8637..72c5c7a3 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/4] 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 72c5c7a3..55296c82 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/4] 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 55296c82..7a2c9f41 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; From 81d8b868e2e38161d475e6e3ca3ff470286e6e49 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 22 Aug 2026 13:55:32 -0400 Subject: [PATCH 4/4] perf(cpu-moe): opt-in AMX tile GEMM for the deduped bf16 pass 1 TDPBF16PS computes C[16x16 fp32] += A[16x32 bf16] x B[32x16 bf16]. Only one mapping of the expert GEMV onto it works: the obvious one (tokens as A's rows) needs B = weights as [K, N] and the bank is [N, K], with no transposed-B form of the instruction. So A is 16 consecutive output rows -- K contiguous, loading straight from the bank at stride H -- B is the routed tokens' activations transposed and VNNI-interleaved, and C comes out as [16 rows x N tokens]. N is the token count, so this only exists on top of the expert dedup: at batch size 1 fifteen of the tile's sixteen columns would be empty. It is off by default (FREETOKEN_CPU_MOE_AMX=1), because on this machine it is worth approximately nothing, and the measurements say precisely why. Deduped bf16 pass 1, E=256 (realistic MoE decode, 2.40 routes/expert at bs=64): bs amx off amx on delta 8 6.39ms 6.29ms +1.5% 16 11.57ms 11.57ms 0.0% 32 17.11ms 17.17ms -0.3% 64 23.80ms 23.11ms +3.0% That is mostly an empty-tile artifact -- with 256 experts the average expert sees 2.4 tokens, so a 16-wide tile runs at ~15% occupancy. Re-running with E=32, where bs=64 gives exactly 16 routes per expert and the tiles fill completely: bs reuse amx off amx on delta 32 8.00x 3.53ms 3.36ms +5.2% 64 16.00x 4.48ms 4.13ms +8.4% So with perfect tile occupancy AMX is worth 8.4%, which is the answer the ISA sweep already gave: avx2 -> avx512f -> avx512bf16 moves 60.7 -> 64.9 -> 67.0 GB/s on this part, so ~10% is all that *any* arithmetic improvement can buy against the memory wall. AMX delivers what is available and not a byte more. Two things follow for anyone reviving this. The tile width is mismatched with MoE decode: filling 16 columns needs ~16 tokens routed to one expert, which at E=256 and top_k=8 means a batch around 512 -- prefill scale, not decode. And the ceiling is a property of the machine: on a part with more DRAM bandwidth per core the same tiles would have more headroom to claim. Correctness runs under the existing GPU-comparison tests with AMX forced on. AMX tile state is an extended XSAVE feature, so it also requests ARCH_REQ_XCOMP_PERM once per process; a refused request leaves amx_ok false and the scalar path in place. FREETOKEN_CPU_MOE_AMX_DEBUG=1 reports whether it actually engaged, which is worth having -- an earlier toggle in this series silently disabled itself and read as "the optimization does nothing". Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01GpGe2fQ5pDShGrnSuksfun --- .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 166 +++++++++++++++++- 1 file changed, 165 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 7a2c9f41..168ab4bf 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -655,6 +655,102 @@ static void cumemop_sync(uintptr_t stream, uintptr_t done_addr, int64_t slot) { "cuStreamWaitValue64(done)"); } +// ================================ AMX (bf16) ================================ +// TDPBF16PS is a tile GEMM: C[16x16 fp32] += A[16x32 bf16] x B[32x16 bf16]. Mapping +// the expert GEMV onto it only works one way round. The obvious assignment -- tokens +// as A's rows -- needs B = weights as [K, N], and the bank is [N, K]; there is no +// transposed-B form of the instruction. So: +// +// A = 16 consecutive *output rows* of the expert (K contiguous -> loads straight +// from the bank, stride H*2) +// B = the routed tokens' activations, transposed and VNNI-interleaved +// C = [16 output rows x N tokens] +// +// N is the token count, so this needs several tokens routed to the same expert -- +// i.e. it rides on the expert dedup, and does nothing at batch size 1 where 15 of +// the tile's 16 columns would be empty. +// +// Expect little on a DRAM-bound part. AMX changes arithmetic throughput, and both +// this and the deduped GEMV read each weight row from memory exactly once; an ISA +// sweep on this machine moves only ~7% from avx2 to avx512bf16, which bounds what +// any arithmetic win can be worth here. Kept opt-in for that reason. +#if defined(__GNUC__) && (__GNUC__ >= 11) && CPU_MOE_X86 +#define CPU_MOE_HAS_AMX 1 +#endif + +#if defined(CPU_MOE_HAS_AMX) +#include + +#define CPU_MOE_ARCH_REQ_XCOMP_PERM 0x1023 +#define CPU_MOE_XFEATURE_XTILEDATA 18 + +struct alignas(64) AmxTileConfig { + uint8_t palette_id; + uint8_t start_row; + uint8_t reserved[14]; + uint16_t colsb[16]; + uint8_t rows[16]; +}; + +// AMX tile state is an extended XSAVE feature the kernel hands out on request; without +// this every tile instruction faults. Process-wide and one-shot. +inline bool amx_request_tile_permission() { + static const bool ok = [] { + return syscall(SYS_arch_prctl, CPU_MOE_ARCH_REQ_XCOMP_PERM, + CPU_MOE_XFEATURE_XTILEDATA) == 0; + }(); + return ok; +} + +inline bool cpu_has_amx_bf16() { + const char* on = getenv("FREETOKEN_CPU_MOE_AMX"); + if (on == nullptr || on[0] == '0') return false; // opt-in, see the note above + if (!__builtin_cpu_supports("amx-tile") || !__builtin_cpu_supports("amx-bf16")) + return false; + return amx_request_tile_permission(); +} + +// tmm0 = C (16x16 fp32), tmm1 = A (16 rows x 32 bf16), tmm2 = B (16 rows x 32 bf16) +__attribute__((target("amx-tile"))) +inline void amx_configure_tiles() { + AmxTileConfig cfg{}; + cfg.palette_id = 1; + for (int t = 0; t < 3; ++t) { + cfg.rows[t] = 16; + cfg.colsb[t] = 64; + } + _tile_loadconfig(&cfg); +} + +// Transpose N token activation rows into the VNNI-interleaved B layout TDPBF16PS +// wants: for K-pair p, row p holds {x[n][2p], x[n][2p+1]} for n = 0..15, so one tile +// row is 32 bf16 = 64 bytes. Columns past N are zeroed and contribute nothing. +inline void amx_pack_b(const bf16_t* const* xs, int n, int K, bf16_t* bv) { + const int kp = K / 2; + std::memset(bv, 0, (size_t)kp * 32 * sizeof(bf16_t)); + for (int p = 0; p < kp; ++p) { + bf16_t* row = bv + (size_t)p * 32; + for (int j = 0; j < n; ++j) { + row[j * 2 + 0] = xs[j][2 * p]; + row[j * 2 + 1] = xs[j][2 * p + 1]; + } + } +} + +// C[16 rows x n tokens] = W[16 rows, K] . B, accumulated in fp32. `ldw` is the bank's +// row stride in bf16 elements; K must be a multiple of 32 (the tile's K depth). +__attribute__((target("amx-tile,amx-bf16"))) +inline void amx_gemm16(const bf16_t* w, int ldw, const bf16_t* bv, int K, float* c16x16) { + _tile_zero(0); + for (int k = 0; k < K; k += 32) { + _tile_loadd(1, w + k, (long)ldw * sizeof(bf16_t)); + _tile_loadd(2, bv + (size_t)(k / 2) * 32, 64); + _tile_dpbf16ps(0, 1, 2); + } + _tile_stored(0, c16x16, 64); +} +#endif // CPU_MOE_HAS_AMX + struct DotChoice { dot_fn fn; const char* name; @@ -1314,6 +1410,8 @@ struct CpuMoeExecutor { // first value seen, which silently disables the A/B when both settings are // exercised from one process. bool dedup_enabled = true; + bool amx_ok = false; // AMX tiles usable (opt-in + CPU + kernel permission) + static constexpr int AMX_MIN_ROUTES = 4; // below this the 16-wide tile is mostly empty std::atomic done_count{0}; std::atomic bar_count{0}; std::atomic bar_sense{0}; @@ -1408,6 +1506,16 @@ struct CpuMoeExecutor { // / 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'; +#if defined(CPU_MOE_HAS_AMX) + // bf16 only: TDPBF16PS eats bf16 tiles, and the quantized formats would have to + // dequantize into a staging tile first, which is the opposite of the point. + amx_ok = (weight_format == WF_BF16) && (H % 32 == 0) && cpu_has_amx_bf16(); + if (getenv("FREETOKEN_CPU_MOE_AMX_DEBUG")) + fprintf(stderr, "[amx] enabled=%d fmt=%d H=%d cpu_tile=%d cpu_bf16=%d perm=%d\n", + (int)amx_ok, weight_format, H, (int)__builtin_cpu_supports("amx-tile"), + (int)__builtin_cpu_supports("amx-bf16"), + (int)amx_request_tile_permission()); +#endif 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") : ""; @@ -1640,6 +1748,56 @@ struct CpuMoeExecutor { // 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. +#if defined(CPU_MOE_HAS_AMX) + // AMX form of the deduped pass 1: 16 output rows x up to 16 routed tokens per tile. + // Falls back to the caller's scalar loop for the row/route remainders. + void do_pass1_amx(const MoeTask* t, int e, int r0, int r1, int i0, int i1) { + thread_local bool cfg_loaded = false; + if (!cfg_loaded) { // TILECFG is per-thread state + amx_configure_tiles(); + cfg_loaded = true; + } + const bf16_t* gate_up_l = reinterpret_cast(tbl_at(gate_up_tbl, t->layer_id)); + thread_local std::vector bv; + bv.resize((size_t)(H / 2) * 32); + const bool swigluoai = act == ACT_SWIGLUOAI; + const float lim = swiglu_limit, alpha = swiglu_alpha; + const bf16_t* xs[16]; + alignas(64) float cg[16 * 16], cu[16 * 16]; + + for (int rc = r0; rc < r1; rc += 16) { + const int n = std::min(16, r1 - rc); + for (int j = 0; j < n; ++j) + xs[j] = t->x + (size_t)(dd_route[rc + j] / top_k) * H; + amx_pack_b(xs, n, H, bv.data()); + for (int i = i0; i + 16 <= i1; i += 16) { + const bf16_t* wg = gate_up_l + ((size_t)e * (2 * I) + i) * H; + const bf16_t* wu = gate_up_l + ((size_t)e * (2 * I) + I + i) * H; + amx_gemm16(wg, H, bv.data(), H, cg); + amx_gemm16(wu, H, bv.data(), H, cu); + for (int rr = 0; rr < 16; ++rr) { + for (int j = 0; j < n; ++j) { + const int route = dd_route[rc + j]; + const float w_in = apply_on_input ? t->w[route] : 1.0f; + float gate = cg[rr * 16 + j] * w_in; + float up = cu[rr * 16 + j] * 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 + rr] = f32_to_bf16(glu * (up + 1.0f)); + } else { + g_row[i + rr] = f32_to_bf16(act_apply(act, gate) * up); + } + } + } + } + } + } +#endif + 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); @@ -1650,10 +1808,16 @@ struct CpuMoeExecutor { 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; + 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; +#if defined(CPU_MOE_HAS_AMX) + if (amx_ok && (r1 - r0) >= AMX_MIN_ROUTES && i0 + 16 <= i1) { + do_pass1_amx(t, e, r0, r1, i0, i1); + i0 += ((i1 - i0) / 16) * 16; // AMX took the whole 16-row groups + } +#endif for (int i = i0; i < i1; ++i) { for (int r = r0; r < r1; ++r) { const int route = dd_route[r];