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..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; @@ -1095,6 +1191,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 +1393,25 @@ 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; + 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}; @@ -1386,6 +1505,17 @@ 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'; +#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") : ""; @@ -1571,13 +1701,165 @@ 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. +#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); + 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)); + 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]; + 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 +1908,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 +2014,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 +2343,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 +2351,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);