From 51afda369c0cada083f549e7953197faf5de211c Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Tue, 13 Oct 2026 08:01:12 +0000 Subject: [PATCH 1/4] refactor(cpu-moe): split the executor into weight formats and gpu sync backends - move each expert weight format into kernels/.cpp behind the MoeFormat interface in format.h; the executor keeps the threads, the work split and the pass 2 route sum - pass each format its banks and integer options by name; MoeParams keeps only the shared shape and activation fields - key the formats by name in one table that weight_formats() exports, replacing the integer format ids and the fp8_block_isa_name hook - replace the per-(layer, batch size) flag slots, pinned IO buffers and task table with one doorbell word MoeTask{layer, num_tokens} and one shared pinned IO block of max_tokens rows - move the GPU/CPU handshake behind GpuSync with a CUDA and a HIP backend picked by setup.py; HIP waits on signal memory and adds explicit batch-memop graph nodes where the runtime does not record the native memops under capture - raise on graph capture when the handshake is off and host functions are not replay-safe - allocate format scratch on 64-byte boundaries - size max_tokens by the explicit cuda_graph_bs as well Co-authored-by: zihaomu --- docs/install_amd.md | 8 + python/freetoken/engine/engine.py | 13 +- python/freetoken/kernel/csrc/cpu_moe/common.h | 235 ++ .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 2280 ++--------------- python/freetoken/kernel/csrc/cpu_moe/format.h | 131 + .../freetoken/kernel/csrc/cpu_moe/gpu_sync.h | 72 + .../kernel/csrc/cpu_moe/gpu_sync_cuda.cpp | 156 ++ .../kernel/csrc/cpu_moe/gpu_sync_hip.cpp | 260 ++ .../kernel/csrc/cpu_moe/kernels/bf16.cpp | 159 ++ .../kernel/csrc/cpu_moe/kernels/ds_fp4.cpp | 269 ++ .../kernel/csrc/cpu_moe/kernels/fp8_block.cpp | 203 ++ .../kernel/csrc/cpu_moe/kernels/mxfp4.cpp | 247 ++ .../kernel/csrc/cpu_moe/kernels/nvfp4.cpp | 466 ++++ .../kernel/csrc/cpu_moe/kernels/q4_0.cpp | 216 ++ python/freetoken/moe/benchbw.py | 9 +- python/freetoken/moe/cpu_executor.py | 439 ++-- setup.py | 5 + tests/moe/test_cpu_moe.py | 10 +- tests/moe/test_cpu_moe_fp8_block.py | 25 +- tests/moe/test_cpu_moe_q4_0.py | 2 +- tests/moe/test_dsfp4_prequant.py | 10 +- 21 files changed, 2837 insertions(+), 2378 deletions(-) create mode 100644 python/freetoken/kernel/csrc/cpu_moe/common.h create mode 100644 python/freetoken/kernel/csrc/cpu_moe/format.h create mode 100644 python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h create mode 100644 python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp create mode 100644 python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp diff --git a/docs/install_amd.md b/docs/install_amd.md index 81d1ad5009..ed934ad657 100644 --- a/docs/install_amd.md +++ b/docs/install_amd.md @@ -45,3 +45,11 @@ On ROCm, these paths use the built-in fallbacks; do not install the `[accel]`, `[fi]`, or `[sgl]` extras. Forcing `--attention-backend fi`, `fa`, or `trtllm`, or NVFP4 Marlin/b12x (e.g. `--quant-backend moe.nvfp4=marlin` or `--quant-backend moe.nvfp4=b12x`), fails with a ROCm-specific error. + +## CPU and hybrid MoE + +`--moe-strategy cpu`, `hybrid` and `--moe-cpu-layers` sync the GPU with the CPU +expert pool through HIP stream memory operations. ROCm 10 records them into CUDA +graphs itself; on ROCm 7.14, which does not, FreeToken adds the graph nodes +explicitly. When the startup probe of that path fails, CUDA graph capture stops +with an error; run those strategies with `--cuda-graph-max-bs 0`. diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 3988caf602..feb33fbb94 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -813,9 +813,8 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: """Build the persistent CPU MoE executor (decode-time expert compute). Must run before CUDA graph capture: the worker pool has to be live for the - eager warmup forward, and the pinned IO buffers / host-func task pointers - must be stable for the captured nodes. Buffers/tasks themselves are - allocated lazily on the first (eager) forward at each batch size. + eager warmup forward, and the pinned IO block must be stable for the captured + nodes. """ from freetoken.moe.cpu_executor import CpuMoeExecutor @@ -828,7 +827,7 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: ) # Decode batches never exceed max_running_req, but CUDA-graph padding can # round a batch up to the largest captured size; cover both. - max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, 1) + max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, *(config.cuda_graph_bs or []), 1) method = sample.quant_method executor = CpuMoeExecutor( cache, @@ -1365,10 +1364,10 @@ def _cpu_moe_executor_viable(model_config) -> bool: """Whether an automatic CPU-decode decision may target the CPU MoE executor. A default boot must degrade to GPU offload instead of crashing in CpuMoeExecutor after the whole load; explicit cpu/hybrid/--moe-cpu-layers picks still fail loudly.""" - from freetoken.moe.cpu_executor import _WFMT_IDS, compiled_extension_supports + from freetoken.moe.cpu_executor import compiled_extension_supports try: - from freetoken.kernel import _cpu_moe # noqa: F401 + from freetoken.kernel import _cpu_moe except ImportError: return False act = getattr(model_config, "hidden_act", "silu") @@ -1379,7 +1378,7 @@ def _cpu_moe_executor_viable(model_config) -> bool: return False expert_quant = getattr(model_config, "expert_quant", "none") fmt = expert_quant if expert_quant != "none" else (moe_wfmt or "bf16") - return fmt == "mxfp4" or fmt in _WFMT_IDS + return fmt == "mxfp4" or fmt in _cpu_moe.weight_formats() def _pin_budget_bytes(reserved: int = 0) -> int | None: diff --git a/python/freetoken/kernel/csrc/cpu_moe/common.h b/python/freetoken/kernel/csrc/cpu_moe/common.h new file mode 100644 index 0000000000..adf7e14584 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/common.h @@ -0,0 +1,235 @@ +// Numeric helpers and ISA detection shared by the CPU MoE formats. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#if defined(__x86_64__) || defined(__i386__) +#include +#define CPU_MOE_X86 1 +#else +#define CPU_MOE_X86 0 +#endif + +// Each SIMD kernel is compiled for its own ISA and picked at runtime (pick_isa), so the build sets no global -march. +#define CPU_MOE_TARGET(...) __attribute__((target(__VA_ARGS__))) + +#if CPU_MOE_X86 && ((defined(__GNUC__) && __GNUC__ >= 10) || defined(__clang__)) +#define CPU_MOE_HAS_AVX512BF16 1 +#define CPU_MOE_HAS_AVX512VNNI 1 +#endif + +namespace freetoken::cpu_moe { + +// Scratch rows are read with 64-byte SIMD loads, so their base is cache-line aligned whatever the heap state; std::vector alone guarantees 16 bytes. +template +struct CacheLineAllocator { + using value_type = T; + CacheLineAllocator() = default; + template + CacheLineAllocator(const CacheLineAllocator&) {} + T* allocate(size_t n) { return static_cast(::operator new(n * sizeof(T), std::align_val_t(64))); } + void deallocate(T* p, size_t) { ::operator delete(p, std::align_val_t(64)); } + template + bool operator==(const CacheLineAllocator&) const { return true; } +}; + +template +using AlignedVector = std::vector>; + +using bf16_t = uint16_t; + +inline float bf16_to_f32(bf16_t v) { + uint32_t u = static_cast(v) << 16; + float f; + std::memcpy(&f, &u, sizeof(f)); + return f; +} + +inline bf16_t f32_to_bf16(float f) { + uint32_t u; + std::memcpy(&u, &f, sizeof(u)); + // round-to-nearest-even + const uint32_t lsb = (u >> 16) & 1u; + u += 0x7fffu + lsb; + return static_cast(u >> 16); +} + +// ACT_SWIGLUOAI is the clamped (up + 1) swiglu (gpt-oss "swigluoai" / +// MiniMax-M3): gate/up are combined jointly with the runtime alpha/limit +// scalars, so it is handled in the gated_rows epilogue (act_apply never sees it; +// the mxfp4 kernel additionally fuses its own copy of the same math). +// ACT_SWIGLU_CLAMP (GLM-5.3 "swiglu_limit") is the same clamped form WITHOUT +// the (up + 1) bias: clamp(gate, max=lim) * sigmoid(alpha*gate) * clamp(up, +-lim). +enum ActKind { + ACT_SILU = 0, + ACT_GELU = 1, + ACT_GELU_TANH = 2, + ACT_SWIGLUOAI = 3, + ACT_SWIGLU_CLAMP = 4, +}; + +inline float act_apply(int act, float x) { + if (act == ACT_SILU) return x / (1.0f + std::exp(-x)); + if (act == ACT_GELU) + return 0.5f * x * (1.0f + std::erf(x * 0.70710678118654752440f)); + // gelu_tanh + const float k0 = 0.79788456080286535588f; // sqrt(2/pi) + const float inner = k0 * (x + 0.044715f * x * x * x); + return 0.5f * x * (1.0f + std::tanh(inner)); +} + +// Software prefetch distance (bytes) ahead of the current weight row stream. The +// weight stream is the bandwidth bottleneck (read once, never reused); nudging the +// HW prefetcher with a few cache lines of lookahead raises sustained throughput. +constexpr int PF_AHEAD = 512; + +const float kE2M1[16] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f}; + +inline float fp16_to_f32(uint16_t h) { + const uint32_t sign = (uint32_t)(h & 0x8000u) << 16; + uint32_t exp = (h >> 10) & 0x1Fu; + uint32_t man = h & 0x3FFu; + uint32_t f; + if (exp == 0) { + if (man == 0) { + f = sign; + } else { + exp = 127 - 15 + 1; + while ((man & 0x400u) == 0) { + man <<= 1; + --exp; + } + man &= 0x3FFu; + f = sign | (exp << 23) | (man << 13); + } + } else if (exp == 0x1Fu) { + f = sign | 0x7F800000u | (man << 13); + } else { + f = sign | ((exp + (127 - 15)) << 23) | (man << 13); + } + float out; + std::memcpy(&out, &f, sizeof(out)); + return out; +} + +// e4m3 (OCP "fn": finite, max-normal 448, exp bias 7). Decoded into a 256-entry LUT. +inline float e4m3_decode(uint8_t v) { + const float sign = (v & 0x80u) ? -1.0f : 1.0f; + const uint32_t exp = (v >> 3) & 0xFu; + const uint32_t man = v & 0x7u; + if (exp == 0) return sign * (man / 8.0f) * 0.015625f; // 2^(1-7) = 2^-6 + return sign * (1.0f + man / 8.0f) * std::ldexp(1.0f, (int)exp - 7); +} + +#if CPU_MOE_X86 +CPU_MOE_TARGET("avx2,fma") +inline float hsum256(__m256 v) { + __m128 lo = _mm256_castps256_ps128(v); + lo = _mm_add_ps(lo, _mm256_extractf128_ps(v, 1)); + lo = _mm_add_ps(lo, _mm_movehl_ps(lo, lo)); + lo = _mm_add_ss(lo, _mm_shuffle_ps(lo, lo, 0x55)); + return _mm_cvtss_f32(lo); +} + +// AVX2 e2m1 nibble decode: codes (0..15) in 8 int32 lanes -> fp32. AVX2 vpermps is +// only 8-wide, so instead of a 16-entry LUT we use the e2m1 sign/magnitude symmetry: +// value = (code&8 ? - : +) * mag8[code&7], mag8 = e2m1[0..7]. The sign is bit 3 of +// the code shifted into the fp32 sign bit (bit 31). Bit-identical to the e2m1 LUT. +CPU_MOE_TARGET("avx2,fma") +inline __m256 e2m1_decode8(__m256i codes, __m256 mag8) { + __m256 mag = _mm256_permutevar8x32_ps(mag8, _mm256_and_si256(codes, _mm256_set1_epi32(7))); + __m256i sgn = _mm256_slli_epi32(_mm256_and_si256(codes, _mm256_set1_epi32(8)), 28); + return _mm256_xor_ps(mag, _mm256_castsi256_ps(sgn)); +} +#endif // CPU_MOE_X86 + +// SIMD tiers, ascending. Each format picks the highest tier <= the one chosen by +// pick_isa() that it implements (fp4 and fp8_block formats have no bf16-specific tier, so the +// avx512bf16 tier maps to their avx512 kernel). +enum IsaTier { ISA_SCALAR = 0, ISA_AVX2 = 1, ISA_AVX512 = 2, ISA_AVX512BF16 = 3 }; + +// Best tier the CPU+build supports, optionally capped DOWN by +// FREETOKEN_CPU_MOE_ISA={scalar,avx2,avx512,avx512bf16} (A/B testing on a machine +// that supports more). FREETOKEN_CPU_MOE_SCALAR=1 forces scalar (legacy alias). +inline IsaTier pick_isa() { +#if CPU_MOE_X86 + if (getenv("FREETOKEN_CPU_MOE_SCALAR")) return ISA_SCALAR; + IsaTier best = ISA_SCALAR; + if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma")) best = ISA_AVX2; + if (best >= ISA_AVX2 && __builtin_cpu_supports("avx512f")) best = ISA_AVX512; +#ifdef CPU_MOE_HAS_AVX512BF16 + if (best >= ISA_AVX512 && __builtin_cpu_supports("avx512bf16")) best = ISA_AVX512BF16; +#endif + if (const char* f = getenv("FREETOKEN_CPU_MOE_ISA")) { + IsaTier want = best; + if (!std::strcmp(f, "scalar")) want = ISA_SCALAR; + else if (!std::strcmp(f, "avx2")) want = ISA_AVX2; + else if (!std::strcmp(f, "avx512")) want = ISA_AVX512; + else if (!std::strcmp(f, "avx512bf16")) want = ISA_AVX512BF16; + if (want < best) best = want; // cap downward; never force above hw/build support + } + return best; +#else + return ISA_SCALAR; +#endif +} + +inline const char* tier_name(IsaTier t) { +#if CPU_MOE_X86 +#ifdef CPU_MOE_HAS_AVX512BF16 + if (t >= ISA_AVX512BF16) return "avx512bf16"; +#endif + if (t >= ISA_AVX512) return "avx512f"; + if (t >= ISA_AVX2) return "avx2"; +#endif + (void)t; + return "scalar"; +} + +// AVX-VNNI (VEX-256 VPDPBUSD) availability: Alder/Raptor Lake, Sapphire Rapids+, Zen5. +// Distinct from AVX-512 VNNI. Opt out with FREETOKEN_CPU_MOE_NO_VNNI=1 (A/B the W4A8 path). +inline bool cpu_has_avxvnni() { +#if CPU_MOE_X86 + const char* no = getenv("FREETOKEN_CPU_MOE_NO_VNNI"); + if (no && no[0] && no[0] != '0') return false; // ignore unset/empty/"0" + return __builtin_cpu_supports("avx2") && __builtin_cpu_supports("avxvnni"); +#else + return false; +#endif +} + +// AVX-512 VNNI (512-bit VPDPBUSD): Cascade Lake+, Ice Lake, Sapphire/Emerald Rapids, +// Zen4+. 2x the 256-bit AVX-VNNI width. FREETOKEN_CPU_MOE_NO_AVX512VNNI=1 forces the +// 256-bit path (A/B the two W4A8 kernels on the same box); FREETOKEN_CPU_MOE_NO_VNNI=1 +// still disables the whole W4A8 family (back to the faithful fp32 nvdot), so it is +// honored here too. +inline bool cpu_has_avx512vnni() { +#if CPU_MOE_X86 && defined(CPU_MOE_HAS_AVX512VNNI) + const char* no = getenv("FREETOKEN_CPU_MOE_NO_AVX512VNNI"); + if (no && no[0] && no[0] != '0') return false; + const char* no_vnni = getenv("FREETOKEN_CPU_MOE_NO_VNNI"); + if (no_vnni && no_vnni[0] && no_vnni[0] != '0') return false; + return __builtin_cpu_supports("avx512vnni"); +#else + return false; +#endif +} + +// Split an interleaved bf16 row into fp32 even/odd halves (even[m]=src[2m]). +// bf16->fp32 is exact, so this only reorders -- done once per token/route and +// reused across every output row of the GEMV. +inline void deinterleave_bf16_f32(const bf16_t* src, float* even, float* odd, int K) { + for (int m = 0; m < K / 2; ++m) { + even[m] = bf16_to_f32(src[2 * m]); + odd[m] = bf16_to_f32(src[2 * m + 1]); + } +} + +} // namespace freetoken::cpu_moe 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 6f85971e9c..40cbd6796c 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -1,32 +1,27 @@ -// CPU-compute MoE executor for the "cpu" offload backend. +// CPU-compute MoE executor for the cpu and hybrid MoE strategies, and its Python binding. // // Decode ships activations to the CPU, computes the routed experts here (reading -// the pinned host expert banks at full RAM bandwidth), and ships the results -// back. To keep the whole decode path inside a single CUDA graph we expose -// submit/sync as host nodes via cudaLaunchHostFunc -- the callbacks only touch a -// CPU worker pool + pinned host buffers and never call any CUDA API. +// the pinned host expert banks at full RAM bandwidth), and ships the results back. +// One task (one MoE layer) is in flight at a time: submit() wakes the persistent worker +// pool, the pool runs pass 1 (gate_up) and pass 2 (down), and sync() waits for it. +// The GPU rings a doorbell to submit and waits on it to sync (GpuSync, gpu_sync.h), so the +// whole decode step stays in one CUDA graph; cudaLaunchHostFunc nodes are the fallback. // -// One task is in flight at a time (per MoE layer): submit() wakes the pool, -// sync() blocks the host-func thread until the pool drains. The heavy GEMV runs -// on the persistent worker threads, not the host-func thread. -// -// Weight formats: bf16, NVFP4, MXFP4, ds_fp4, Q4_0 and block-FP8 expert banks (see WFmt and -// the per-format bank schemas). Compute is FP32-accumulate; the intermediate is -// stored bf16 to match the GPU decode path. ISA is chosen once at construction -// (AVX-512-BF16 dpbf16 -> AVX-512F widening -> AVX2+FMA -> scalar). +// Weight formats live in kernels/, one file each behind the MoeFormat interface (format.h). +// Compute is FP32-accumulate; the intermediate is stored bf16 to match the GPU decode path. #include #include #include -#include #include #include -#include #include +#include #include #include #include #include +#include #include #include @@ -40,1379 +35,98 @@ #define CPU_MOE_HAS_AFFINITY 0 #endif -#if defined(__x86_64__) || defined(__i386__) -#include -#define CPU_MOE_X86 1 -#else -#define CPU_MOE_X86 0 -#endif - -namespace { - -inline void cpu_relax() { -#if CPU_MOE_X86 - _mm_pause(); -#elif defined(__aarch64__) - asm volatile("yield" ::: "memory"); -#endif -} - -using bf16_t = uint16_t; - -inline float bf16_to_f32(bf16_t v) { - uint32_t u = static_cast(v) << 16; - float f; - std::memcpy(&f, &u, sizeof(f)); - return f; -} - -inline bf16_t f32_to_bf16(float f) { - uint32_t u; - std::memcpy(&u, &f, sizeof(u)); - // round-to-nearest-even - const uint32_t lsb = (u >> 16) & 1u; - u += 0x7fffu + lsb; - return static_cast(u >> 16); -} - -// ACT_SWIGLUOAI is the clamped (up + 1) swiglu (gpt-oss "swigluoai" / -// MiniMax-M3): gate/up are combined jointly with the runtime alpha/limit -// scalars, so it is handled in the do_pass1 epilogue (act_apply never sees it; -// the mxfp4 kernel additionally fuses its own copy of the same math). -// ACT_SWIGLU_CLAMP (GLM-5.3 "swiglu_limit") is the same clamped form WITHOUT -// the (up + 1) bias: clamp(gate, max=lim) * sigmoid(alpha*gate) * clamp(up, +-lim). -enum ActKind { - ACT_SILU = 0, - ACT_GELU = 1, - ACT_GELU_TANH = 2, - ACT_SWIGLUOAI = 3, - ACT_SWIGLU_CLAMP = 4, -}; - -inline float act_apply(int act, float x) { - if (act == ACT_SILU) return x / (1.0f + std::exp(-x)); - if (act == ACT_GELU) - return 0.5f * x * (1.0f + std::erf(x * 0.70710678118654752440f)); - // gelu_tanh - const float k0 = 0.79788456080286535588f; // sqrt(2/pi) - const float inner = k0 * (x + 0.044715f * x * x * x); - return 0.5f * x * (1.0f + std::tanh(inner)); -} - -// ------------------------------- dot products ------------------------------- -// dot(weight[bf16], act[bf16], n) -> fp32. The selected impl is a function -// pointer chosen at runtime; the per-row call overhead is negligible vs n. - -using dot_fn = float (*)(const bf16_t*, const bf16_t*, int); - -float dot_scalar(const bf16_t* w, const bf16_t* x, int n) { - float acc = 0.0f; - for (int i = 0; i < n; ++i) acc += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); - return acc; -} - -// Software prefetch distance (bytes) ahead of the current weight row stream. The -// weight stream is the bandwidth bottleneck (read once, never reused); nudging the -// HW prefetcher with a few cache lines of lookahead raises sustained throughput. -constexpr int PF_AHEAD = 512; - -#if CPU_MOE_X86 -__attribute__((target("avx512f"))) -float dot_avx512f(const bf16_t* w, const bf16_t* x, int n) { - // 4 independent accumulators -> more in-flight loads (memory-level parallelism), - // which is what lifts a bandwidth-bound GEMV toward peak. - __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); - __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); - int i = 0; - for (; i + 64 <= n; i += 64) { - _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); - for (int j = 0; j < 64; j += 16) { - __m256i wi = _mm256_loadu_si256(reinterpret_cast(w + i + j)); - __m256i xi = _mm256_loadu_si256(reinterpret_cast(x + i + j)); - __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16)); - __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16)); - __m512& acc = (j == 0) ? a0 : (j == 16) ? a1 : (j == 32) ? a2 : a3; - acc = _mm512_fmadd_ps(wf, xf, acc); - } - } - for (; i + 16 <= n; i += 16) { - __m256i wi = _mm256_loadu_si256(reinterpret_cast(w + i)); - __m256i xi = _mm256_loadu_si256(reinterpret_cast(x + i)); - __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16)); - __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16)); - a0 = _mm512_fmadd_ps(wf, xf, a0); - } - float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3))); - for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); - return s; -} - -#if (defined(__GNUC__) && __GNUC__ >= 10) || defined(__clang__) -#define CPU_MOE_HAS_AVX512BF16 1 -__attribute__((target("avx512bf16,avx512f"))) -static inline __m512bh load_bh(const bf16_t* p) { - __m512i raw = _mm512_loadu_si512(reinterpret_cast(p)); - __m512bh out; - std::memcpy(&out, &raw, sizeof(out)); - return out; -} - -__attribute__((target("avx512bf16,avx512f"))) -float dot_avx512bf16(const bf16_t* w, const bf16_t* x, int n) { - // 4 accumulators (128 bf16/iter) for memory-level parallelism + a prefetch nudge. - __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); - __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); - int i = 0; - for (; i + 128 <= n; i += 128) { - _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); - a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i)); - a1 = _mm512_dpbf16_ps(a1, load_bh(w + i + 32), load_bh(x + i + 32)); - a2 = _mm512_dpbf16_ps(a2, load_bh(w + i + 64), load_bh(x + i + 64)); - a3 = _mm512_dpbf16_ps(a3, load_bh(w + i + 96), load_bh(x + i + 96)); - } - for (; i + 32 <= n; i += 32) { - a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i)); - } - float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3))); - for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); - return s; -} -#endif // avx512bf16 available - -// --------- AVX2 + FMA (256-bit fallback for CPUs without AVX-512) ---------- -// Covers Intel 12-14th gen / Arrow Lake (AVX-512 fused off) and AMD Zen<4. The -// bf16->fp32 widen is a zero-extend + <<16; with 4 independent accumulators the -// GEMV is memory-bandwidth bound, same as the AVX-512 path (just half the width). -__attribute__((target("avx2,fma"))) -inline float hsum256(__m256 v) { - __m128 lo = _mm256_castps256_ps128(v); - lo = _mm_add_ps(lo, _mm256_extractf128_ps(v, 1)); - lo = _mm_add_ps(lo, _mm_movehl_ps(lo, lo)); - lo = _mm_add_ss(lo, _mm_shuffle_ps(lo, lo, 0x55)); - return _mm_cvtss_f32(lo); -} - -__attribute__((target("avx2,fma"))) -float dot_avx2(const bf16_t* w, const bf16_t* x, int n) { - __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps(); - __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps(); - int i = 0; - for (; i + 32 <= n; i += 32) { - _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); - for (int j = 0; j < 32; j += 8) { - __m128i wi = _mm_loadu_si128(reinterpret_cast(w + i + j)); - __m128i xi = _mm_loadu_si128(reinterpret_cast(x + i + j)); - __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16)); - __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16)); - __m256& acc = (j == 0) ? a0 : (j == 8) ? a1 : (j == 16) ? a2 : a3; - acc = _mm256_fmadd_ps(wf, xf, acc); - } - } - for (; i + 8 <= n; i += 8) { - __m128i wi = _mm_loadu_si128(reinterpret_cast(w + i)); - __m128i xi = _mm_loadu_si128(reinterpret_cast(x + i)); - __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16)); - __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16)); - a0 = _mm256_fmadd_ps(wf, xf, a0); - } - float s = hsum256(_mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3))); - for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); - return s; -} -#endif // CPU_MOE_X86 +#include "format.h" +#include "gpu_sync.h" -// --------------------------- NVFP4 (W4A16) dequant --------------------------- -// Weights: e2m1 4-bit codes (2/byte, low nibble first), per-16 block scale in -// fp8-e4m3, per-output-row global scale in fp16. Dequant matches the GPU kernels -// (freetoken/kernel/triton/nvfp4_dequant.py): w = E2M1[code] * e4m3(scale) * global. -// Activations stay bf16 (W4A16); the GEMV dequantizes weights inside the K-loop. +namespace freetoken::cpu_moe { -const float kE2M1[16] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, - -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f}; +std::unique_ptr make_bf16_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_nvfp4_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_mxfp4_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_ds_fp4_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_q4_0_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_fp8_block_format(const MoeParams& p, const FormatArgs& a); -// e2m1 * 2 as exact int8 (all e2m1 values are multiples of 0.5). Used by the AVX-VNNI -// W4A8 path: nibble -> int8 weight via PSHUFB LUT, then VPDPBUSD against int8 activations; -// the *2 is undone by a 0.5 folded into the final scale. Mirrors ggml's kvalues_mxfp4. -alignas(16) const int8_t kE2M1x2[16] = {0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12}; - -inline float fp16_to_f32(uint16_t h) { - const uint32_t sign = (uint32_t)(h & 0x8000u) << 16; - uint32_t exp = (h >> 10) & 0x1Fu; - uint32_t man = h & 0x3FFu; - uint32_t f; - if (exp == 0) { - if (man == 0) { - f = sign; - } else { - exp = 127 - 15 + 1; - while ((man & 0x400u) == 0) { - man <<= 1; - --exp; - } - man &= 0x3FFu; - f = sign | (exp << 23) | (man << 13); - } - } else if (exp == 0x1Fu) { - f = sign | 0x7F800000u | (man << 13); - } else { - f = sign | ((exp + (127 - 15)) << 23) | (man << 13); - } - float out; - std::memcpy(&out, &f, sizeof(out)); - return out; -} - -// e4m3 (OCP "fn": finite, max-normal 448, exp bias 7). Decoded into a 256-entry LUT. -inline float e4m3_decode(uint8_t v) { - const float sign = (v & 0x80u) ? -1.0f : 1.0f; - const uint32_t exp = (v >> 3) & 0xFu; - const uint32_t man = v & 0x7u; - if (exp == 0) return sign * (man / 8.0f) * 0.015625f; // 2^(1-7) = 2^-6 - return sign * (1.0f + man / 8.0f) * std::ldexp(1.0f, (int)exp - 7); -} - -// ------------------------ block-FP8 (W8A16) dequant ------------------------ -// fp8-e4m3 weights row-major [rows, K] plus one bf16 scale per 128x128 block. The scale tensor is named weight_scale_inv but multiplies, as in kernel/triton/fp8_blockscale_moe.py. -// Whole 128-blocks read the activations as fp32 even/odd halves (xe[m] = x[2m], xo[m] = x[2m+1]); a ragged tail reads the bf16 row x. -// 0x7F/0xFF (NaN in torch float8_e4m3fn) decode to +-480 in every kernel, as e4m3_decode does. -using fp8_block_dot_fn = float (*)(const uint8_t*, const float*, const float*, const bf16_t*, - const bf16_t*, int, const float*); - -constexpr int FP8_BLK = 128; // K elements covered by one weight scale - -inline float fp8_block_tail(const uint8_t* w, const bf16_t* x, int k0, int K, const float* e4m3) { - float blk = 0.0f; - for (int k = k0; k < K; ++k) blk += e4m3[w[k]] * bf16_to_f32(x[k]); - return blk; -} - -float dot_fp8_block_scalar(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, - const bf16_t* s, int K, const float* e4m3) { - float acc = 0.0f; - int b = 0, k0 = 0; - for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { - // 8 independent sums: a single running sum is bound by the add latency - float a[8] = {}; - const uint8_t* wb = w + k0; - const float* pe = xe + k0 / 2; - const float* po = xo + k0 / 2; - for (int j = 0; j < FP8_BLK / 2; j += 4) - for (int l = 0; l < 4; ++l) { - a[2 * l] += e4m3[wb[2 * (j + l)]] * pe[j + l]; - a[2 * l + 1] += e4m3[wb[2 * (j + l) + 1]] * po[j + l]; - } - acc += (((a[0] + a[1]) + (a[2] + a[3])) + ((a[4] + a[5]) + (a[6] + a[7]))) * - bf16_to_f32(s[b]); - } - if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); - return acc; -} - -#if CPU_MOE_X86 -// A byte b read as fp16 bits (b & 0x80) << 8 | (b & 0x7F) << 7 is its e4m3 value times 2^-8, subnormals included, so the block scale carries the 256. -// In a 16-bit lane holding bytes (2m, 2m+1), srai 1 then & 0xBF80 builds that for byte 2m+1; a left shift by 8 first does byte 2m. Every AVX2 CPU has F16C. -__attribute__((target("avx2,fma,f16c"))) -float dot_fp8_block_avx2(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, - const bf16_t* s, int K, const float* e4m3) { - const __m256i keep = _mm256_set1_epi16((short)0xBF80); - __m256 total = _mm256_setzero_ps(); - int b = 0, k0 = 0; - for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { - _mm_prefetch(reinterpret_cast(w + k0) + PF_AHEAD, _MM_HINT_T0); - __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps(); - __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps(); - for (int j = 0; j < FP8_BLK; j += 32) { - const __m256i v = _mm256_loadu_si256(reinterpret_cast(w + k0 + j)); - const __m256i odd = _mm256_and_si256(_mm256_srai_epi16(v, 1), keep); - const __m256i even = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v, 8), 1), keep); - const float* pe = xe + (k0 + j) / 2; - const float* po = xo + (k0 + j) / 2; - a0 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(even)), _mm256_loadu_ps(pe), a0); - a1 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(even, 1)), _mm256_loadu_ps(pe + 8), a1); - a2 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(odd)), _mm256_loadu_ps(po), a2); - a3 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(odd, 1)), _mm256_loadu_ps(po + 8), a3); - } - const __m256 blk = _mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3)); - total = _mm256_fmadd_ps(blk, _mm256_set1_ps(bf16_to_f32(s[b]) * 256.0f), total); - } - float acc = hsum256(total); - if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); - return acc; -} - -__attribute__((target("avx512f,avx2,fma,f16c"))) -float dot_fp8_block_avx512f(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, - const bf16_t* s, int K, const float* e4m3) { - const __m256i keep = _mm256_set1_epi16((short)0xBF80); - __m512 total = _mm512_setzero_ps(); - int b = 0, k0 = 0; - for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { - _mm_prefetch(reinterpret_cast(w + k0) + PF_AHEAD, _MM_HINT_T0); - __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); - __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); - for (int j = 0; j < FP8_BLK; j += 64) { - const __m256i v0 = _mm256_loadu_si256(reinterpret_cast(w + k0 + j)); - const __m256i v1 = _mm256_loadu_si256(reinterpret_cast(w + k0 + j + 32)); - const __m256i o0 = _mm256_and_si256(_mm256_srai_epi16(v0, 1), keep); - const __m256i e0 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v0, 8), 1), keep); - const __m256i o1 = _mm256_and_si256(_mm256_srai_epi16(v1, 1), keep); - const __m256i e1 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v1, 8), 1), keep); - const float* pe = xe + (k0 + j) / 2; - const float* po = xo + (k0 + j) / 2; - a0 = _mm512_fmadd_ps(_mm512_cvtph_ps(e0), _mm512_loadu_ps(pe), a0); - a1 = _mm512_fmadd_ps(_mm512_cvtph_ps(o0), _mm512_loadu_ps(po), a1); - a2 = _mm512_fmadd_ps(_mm512_cvtph_ps(e1), _mm512_loadu_ps(pe + 16), a2); - a3 = _mm512_fmadd_ps(_mm512_cvtph_ps(o1), _mm512_loadu_ps(po + 16), a3); - } - const __m512 blk = _mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3)); - total = _mm512_fmadd_ps(blk, _mm512_set1_ps(bf16_to_f32(s[b]) * 256.0f), total); - } - float acc = _mm512_reduce_add_ps(total); - if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); - return acc; -} -#endif // CPU_MOE_X86 - -// Activations pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]); see the -// ds_fp4 dot below for why (drops the hot loop to ~1.5 shuffle ops / 16 weights). -using nvdot_fn = float (*)(const uint8_t*, const uint8_t*, float, const float*, const float*, - int, const float*, const float*); - -float dot_nvfp4_scalar(const uint8_t* packed, const uint8_t* scale, float global, - const float* xe, const float* xo, int K, const float* e2m1, - const float* e4m3) { - float acc = 0.0f; - const int nb = K / 16; - for (int b = 0; b < nb; ++b) { - const float bs = e4m3[scale[b]]; - const uint8_t* pk = packed + (size_t)b * 8; - const float* xeb = xe + (size_t)b * 8; // 16 K -> 8 even + 8 odd - const float* xob = xo + (size_t)b * 8; - float bsum = 0.0f; - for (int j = 0; j < 8; ++j) { - const uint8_t byte = pk[j]; - bsum += e2m1[byte & 0xF] * xeb[j]; - bsum += e2m1[byte >> 4] * xob[j]; - } - acc += bs * bsum; - } - return acc * global; -} - -// ---- NVFP4 W4A8 (int8 activations) dot: nibble->int8 via LUT, per-16 act scale ---- -// asi8: int8 activations laid out per-16 block as [even(8), odd(8)]; asb[b] = per-block -// activation scale (absmax/127). Result folds the e2m1*2 -> *0.5 into the scale. -using nvi8dot_fn = float (*)(const uint8_t*, const uint8_t*, float, const int8_t*, int, - const float*, const float*); - -[[maybe_unused]] float dot_nvfp4_i8_scalar(const uint8_t* packed, const uint8_t* scale, - float global, const int8_t* asi8, int K, const float* e4m3, - const float* asb) { - float acc = 0.0f; - const int nb = K / 16; - for (int b = 0; b < nb; ++b) { - const uint8_t* pk = packed + (size_t)b * 8; - const int8_t* ae = asi8 + (size_t)b * 16; // even(8) - const int8_t* ao = ae + 8; // odd(8) - int isum = 0; - for (int j = 0; j < 8; ++j) { - isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j]; - isum += (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; - } - acc += (e4m3[scale[b]] * asb[b]) * (float)isum; - } - return acc * (0.5f * global); -} - -#if CPU_MOE_X86 -// AVX2 e2m1 nibble decode: codes (0..15) in 8 int32 lanes -> fp32. AVX2 vpermps is -// only 8-wide, so instead of a 16-entry LUT we use the e2m1 sign/magnitude symmetry: -// value = (code&8 ? - : +) * mag8[code&7], mag8 = e2m1[0..7]. The sign is bit 3 of -// the code shifted into the fp32 sign bit (bit 31). Bit-identical to the e2m1 LUT. -__attribute__((target("avx2,fma"))) -inline __m256 e2m1_decode8(__m256i codes, __m256 mag8) { - __m256 mag = _mm256_permutevar8x32_ps(mag8, _mm256_and_si256(codes, _mm256_set1_epi32(7))); - __m256i sgn = _mm256_slli_epi32(_mm256_and_si256(codes, _mm256_set1_epi32(8)), 28); - return _mm256_xor_ps(mag, _mm256_castsi256_ps(sgn)); -} - -// Two 16-K blocks (16 packed bytes) per iter: lo nibbles -> even-K, hi -> odd-K, -// gathered via two vpermps. The per-16 e4m3 scale differs across the two blocks, so -// it is applied per lane (low 8 lanes = block b, high 8 = block b+1). -__attribute__((target("avx512f"))) -inline __m512 nvfp4_blk2(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut, - __m512i loma, float s0, float s1) { - __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast(pk))); - __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); - __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); - __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob))); - // lanes 0-7 (block b) -> s0, lanes 8-15 (block b+1) -> s1 (pure AVX512F mask move) - __m512 scv = _mm512_mask_mov_ps(_mm512_set1_ps(s0), 0xFF00, _mm512_set1_ps(s1)); - return _mm512_mul_ps(prod, scv); -} - -__attribute__((target("avx512f"))) -float dot_nvfp4_avx512(const uint8_t* packed, const uint8_t* scale, float global, - const float* xe, const float* xo, int K, const float* e2m1, - const float* e4m3) { - const __m512 lut = _mm512_loadu_ps(e2m1); // 16 e2m1 values - const __m512i loma = _mm512_set1_epi32(0xF); - __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); - const int nb = K / 16; // 8 packed bytes + one e4m3 scale per block - int b = 0; - for (; b + 4 <= nb; b += 4) { // two blk2 calls -> 4 blocks, 2 accumulators - acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8, - xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]])); - acc1 = _mm512_add_ps(acc1, nvfp4_blk2(packed + (size_t)(b + 2) * 8, xe + (size_t)(b + 2) * 8, - xo + (size_t)(b + 2) * 8, lut, loma, e4m3[scale[b + 2]], e4m3[scale[b + 3]])); - } - for (; b + 2 <= nb; b += 2) - acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8, - xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]])); - float s = _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1)); - for (; b < nb; ++b) { // odd final 16-K block - const uint8_t* pk = packed + (size_t)b * 8; - const float* xeb = xe + (size_t)b * 8; - const float* xob = xo + (size_t)b * 8; - float bsum = 0.0f; - for (int j = 0; j < 8; ++j) { - bsum += e2m1[pk[j] & 0xF] * xeb[j] + e2m1[pk[j] >> 4] * xob[j]; - } - s += e4m3[scale[b]] * bsum; - } - return s * global; -} - -// AVX2: one 16-K block (8 packed bytes) per call, 8 even + 8 odd lanes. -__attribute__((target("avx2,fma"))) -inline __m256 nvfp4_blk_avx2(const uint8_t* pk, const float* xeb, const float* xob, - __m256 mag8, float sc) { - __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); - __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); - __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); - __m256 prod = _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); - return _mm256_mul_ps(prod, _mm256_set1_ps(sc)); -} - -__attribute__((target("avx2,fma"))) -float dot_nvfp4_avx2(const uint8_t* packed, const uint8_t* scale, float global, - const float* xe, const float* xo, int K, const float* e2m1, - const float* e4m3) { - const __m256 mag8 = _mm256_loadu_ps(e2m1); // e2m1[0..7] magnitudes - __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); - const int nb = K / 16; - int b = 0; - for (; b + 2 <= nb; b += 2) { - acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8, - xo + (size_t)b * 8, mag8, e4m3[scale[b]])); - acc1 = _mm256_add_ps(acc1, nvfp4_blk_avx2(packed + (size_t)(b + 1) * 8, xe + (size_t)(b + 1) * 8, - xo + (size_t)(b + 1) * 8, mag8, e4m3[scale[b + 1]])); - } - for (; b < nb; ++b) - acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8, - xo + (size_t)b * 8, mag8, e4m3[scale[b]])); - return hsum256(_mm256_add_ps(acc0, acc1)) * global; -} - -// AVX-VNNI W4A8: decode 8 packed bytes (16 nibbles) of one 16-block to int8 [lo(8),hi(8)] -// via PSHUFB against the e2m1*2 LUT (replaces the 2 vpermps fp32 expands -- ~4x less -// port-5 traffic). lo=even-K weights, hi=odd-K, matching the [even(8),odd(8)] act layout. -__attribute__((target("avx2,avxvnni,fma"))) -inline __m128i nvfp4_decode_block_i8(const uint8_t* pk, __m128i lut) { - __m128i b = _mm_loadl_epi64(reinterpret_cast(pk)); // 8 bytes - __m128i lo = _mm_and_si128(b, _mm_set1_epi8(0x0F)); - __m128i hi = _mm_and_si128(_mm_srli_epi16(b, 4), _mm_set1_epi8(0x0F)); - return _mm_shuffle_epi8(lut, _mm_unpacklo_epi64(lo, hi)); // [lo(8),hi(8)] -> int8 -} - -// Two 16-blocks per VPDPBUSD (32 int8). Sign trick (ggml mul_add_epi8): |w|*(sign(w)*a)=w*a, -// so u8*s8 needs no offset/correction term. Per-block scale (e4m3 * act-scale) folded in fp32 -// (lanes 0-3 -> block b, 4-7 -> block b+1). Bit-faithful weight; only the int8 activation quant -// (W4A8) differs from the bf16 reference. -__attribute__((target("avx2,avxvnni,fma"))) -float dot_nvfp4_i8_vnni(const uint8_t* packed, const uint8_t* scale, float global, - const int8_t* asi8, int K, const float* e4m3, const float* asb) { - const __m128i lut = _mm_loadu_si128(reinterpret_cast(kE2M1x2)); - __m256 accF = _mm256_setzero_ps(); - const int nb = K / 16; - int b = 0; - for (; b + 2 <= nb; b += 2) { - __m128i wb = nvfp4_decode_block_i8(packed + (size_t)b * 8, lut); - __m128i wb1 = nvfp4_decode_block_i8(packed + (size_t)(b + 1) * 8, lut); - __m256i w = _mm256_set_m128i(wb1, wb); // [blk b | blk b+1] - __m256i a = _mm256_loadu_si256(reinterpret_cast(asi8 + (size_t)b * 16)); - __m256i aw = _mm256_sign_epi8(w, w); // |w| (u8 operand) - __m256i sa = _mm256_sign_epi8(a, w); // sign(w)*a (s8 operand) - __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); - __m256 scv = _mm256_blend_ps(_mm256_set1_ps(e4m3[scale[b]] * asb[b]), - _mm256_set1_ps(e4m3[scale[b + 1]] * asb[b + 1]), 0xF0); - accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), scv, accF); - } - float s = hsum256(accF); - for (; b < nb; ++b) { // tail (odd block count) - const uint8_t* pk = packed + (size_t)b * 8; - const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; - int isum = 0; - for (int j = 0; j < 8; ++j) - isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; - s += (e4m3[scale[b]] * asb[b]) * (float)isum; - } - return s * (0.5f * global); -} - -#if (defined(__GNUC__) && __GNUC__ >= 10) || defined(__clang__) -#define CPU_MOE_HAS_AVX512VNNI 1 - -// Software-prefetch distance for the W4A8 weight stream, in 16-K blocks (8 packed -// bytes each). Returns -1 when FREETOKEN_CPU_MOE_PF_BLOCKS is unset: the kernel then -// uses the built-in default min(512 blocks = 4 KB, 2 rows) -- 4 KB is the empirical -// optimum on large-row machines (Emerald Rapids sweep), while the 2-row cap keeps a small-row -// model's overshoot bounded (the executor works in 32-row tiles, so a fixed byte -// distance otherwise prefetches another worker's tile: duplicated DRAM traffic that -// regresses at the bandwidth ceiling). An EXPLICIT env value is honored verbatim -// (no clamp; 0 disables): the per-machine optimum can sit past the safe default (+20% -// at 4 KB on a 24-thread Ice Lake with 256B rows), so the escape hatch must reach it. -// Prefetch never faults, so overshooting a row/bank tail is safe. -static int nvfp4_pf_blocks() { - static const int v = [] { - const char* s = getenv("FREETOKEN_CPU_MOE_PF_BLOCKS"); - return (s && s[0]) ? atoi(s) : -1; - }(); - return v; -} - -// AVX-512 VNNI W4A8: FOUR 16-K blocks per VPDPBUSD (64 int8) -- 2x the AVX-VNNI -// (256-bit) path. Decode 32 packed bytes -> 64 int8 with a single _mm512_shuffle_epi8 -// (e2m1*2 LUT replicated to all 4 128-bit lanes). AVX-512 has no _mm512_sign_epi8, so -// the u8*s8 sign trick (|w| as u8, sign(w)*a as s8) uses abs_epi8 + a masked negate. -// Bit-faithful weight; only the int8 activation quant (W4A8) differs from bf16. -// One 4-block group (64 int8) -> scaled fp32 partial (16 lanes). Isolated as a helper so -// the caller can run several independent chains into separate accumulators (the decode -> -// dpbusd -> scale chain is long, so a single accumulator leaves the core latency-bound). -__attribute__((target("avx512f,avx512bw,avx512vnni,avx2"))) -static inline __m512 nvfp4_i8_grp4(const uint8_t* packed, const uint8_t* scale, - const int8_t* asi8, const float* e4m3, const float* asb, - int b, __m512i lut, __m512i idx, __m512i mask0F, - __m512i idxsc) { - const __mmask64 hi_half = 0xFF00FF00FF00FF00ULL; // bytes 8-15 of each 128b lane -> hi nibbles - __m256i raw = _mm256_loadu_si256(reinterpret_cast(packed + (size_t)b * 8)); - __m512i src = _mm512_permutexvar_epi64(idx, _mm512_castsi256_si512(raw)); - __m512i lo = _mm512_and_si512(src, mask0F); - __m512i hi = _mm512_and_si512(_mm512_srli_epi16(src, 4), mask0F); - __m512i comb = _mm512_mask_blend_epi8(hi_half, lo, hi); - __m512i w = _mm512_shuffle_epi8(lut, comb); // 64 int8 weights (e2m1*2) - // u8*s8 sign trick without _mm512_sign_epi8: aw=|w|, sa = (w<0 ? -a : a) (a moot at w==0) - __m512i a = _mm512_loadu_si512(reinterpret_cast(asi8 + (size_t)b * 16)); - __m512i aw = _mm512_abs_epi8(w); - __mmask64 neg = _mm512_movepi8_mask(w); - __m512i sa = _mm512_mask_sub_epi8(a, neg, _mm512_setzero_si512(), a); - __m512i di = _mm512_dpbusd_epi32(_mm512_setzero_si512(), aw, sa); // 16 int32 (groups of 4) - // int32[0..3]->blk b, [4..7]->b+1, [8..11]->b+2, [12..15]->b+3. The 4 block scales - // (e4m3 LUT x per-block act scale) are computed vectorized: a 4-byte load + epu8->epi32 - // widen + one 4-lane LUT gather + one mul replaces 8 scalar loads + 4 scalar muls + - // a set_ps assembly, which otherwise dominates the per-group op count. - int sc_raw; - memcpy(&sc_raw, scale + b, 4); - __m128i sc4 = _mm_cvtepu8_epi32(_mm_cvtsi32_si128(sc_raw)); - __m128 s4 = _mm_mul_ps(_mm_i32gather_ps(e4m3, sc4, 4), _mm_loadu_ps(asb + b)); - __m512 scv = _mm512_permutexvar_ps(idxsc, _mm512_castps128_ps512(s4)); - return _mm512_mul_ps(_mm512_cvtepi32_ps(di), scv); -} - -__attribute__((target("avx512f,avx512bw,avx512vnni,avx2"))) -float dot_nvfp4_i8_avx512vnni(const uint8_t* packed, const uint8_t* scale, float global, - const int8_t* asi8, int K, const float* e4m3, const float* asb) { - const __m512i lut = _mm512_broadcast_i32x4( - _mm_loadu_si128(reinterpret_cast(kE2M1x2))); - const __m512i idx = _mm512_set_epi64(3, 3, 2, 2, 1, 1, 0, 0); // block i -> 128b lane i - const __m512i mask0F = _mm512_set1_epi8(0x0F); - const __m512i idxsc = _mm512_set_epi32(3, 3, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0); - __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); - __m512 acc2 = _mm512_setzero_ps(), acc3 = _mm512_setzero_ps(); - const int nb = K / 16; - // Built-in default: min(4 KB, 2 rows) -- see nvfp4_pf_blocks(). An explicitly set - // env value is used verbatim so a per-machine sweep can reach operating points past - // the conservative default. - const int pfb = nvfp4_pf_blocks(); - const int pf = (pfb < 0) ? std::min(512, 2 * nb) : pfb; - int b = 0; - // Four independent 4-block groups per iter -> four accumulator chains hide the ~10-op - // decode->dpbusd->scale latency (1-2 chains leave the loop latency-bound: the per-core - // rate sat at ~60% of the core's achievable DRAM stream rate). - for (; b + 16 <= nb; b += 16) { - // Prefetch the weight stream ahead: the interleaved decode lowers the L1-miss - // concurrency the HW prefetcher sustains on its own (134 -> 167 GB/s on Emerald Rapids). - if (pf > 0) { - _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8), - _MM_HINT_T0); - _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8 + 64), - _MM_HINT_T0); - } - acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, - mask0F, idxsc)); - acc1 = _mm512_add_ps(acc1, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 4, lut, idx, - mask0F, idxsc)); - acc2 = _mm512_add_ps(acc2, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 8, lut, idx, - mask0F, idxsc)); - acc3 = _mm512_add_ps(acc3, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 12, lut, idx, - mask0F, idxsc)); - } - for (; b + 4 <= nb; b += 4) - acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, - mask0F, idxsc)); - float s = _mm512_reduce_add_ps( - _mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3))); - for (; b < nb; ++b) { // tail (<4 remaining 16-K blocks) - const uint8_t* pk = packed + (size_t)b * 8; - const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; - int isum = 0; - for (int j = 0; j < 8; ++j) - isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; - s += (e4m3[scale[b]] * asb[b]) * (float)isum; - } - return s * (0.5f * global); -} -#endif // avx512vnni available -#endif - -// ===================================================================================== -// CUDA stream memory operations (driver API, resolved via dlopen -- no link-time or -// toolchain dependence). The GPU side of the flag handshake: submit = WRITE_VALUE -// (done[slot]=0 then ready[slot]=1), sync = WAIT_VALUE(done[slot] >= 1). The wait is -// executed by the GPU front-end (no SM-resident kernel), so GPU "utilization" stays -// truthful during CPU compute windows -- a resident spin kernel pinned it at 99%, -// which laptop CPU/GPU dynamic power schedulers answered by clamping the CPU's max -// frequency (GEMV workers -1.5x: the reported edge regression). Availability is -// probed functionally at startup (memops_probe); anything unsupported (Windows WDDM, -// vGPU, old drivers) falls back to the cudaLaunchHostFunc path. -#if defined(_WIN32) -#include -static void* cumemop_dlopen() { return (void*)::LoadLibraryA("nvcuda.dll"); } -static void* cumemop_dlsym(void* h, const char* n) { - return (void*)::GetProcAddress((HMODULE)h, n); -} -#else -#include -static void* cumemop_dlopen() { - void* h = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL); - if (h == nullptr) h = dlopen("libcuda.so", RTLD_LAZY | RTLD_LOCAL); - return h; -} -static void* cumemop_dlsym(void* h, const char* n) { return dlsym(h, n); } -#endif - -using cuMemOp64_fn = int (*)(void* stream, unsigned long long addr, unsigned long long value, - unsigned int flags); -static cuMemOp64_fn g_cu_write64 = nullptr; -static cuMemOp64_fn g_cu_wait64 = nullptr; -static constexpr unsigned int kCuWaitValueGeq = 0x0; // CU_STREAM_WAIT_VALUE_GEQ -static constexpr unsigned int kCuWriteDefault = 0x0; // CU_STREAM_WRITE_VALUE_DEFAULT - -static bool cumemop_resolve() { - static bool resolved = [] { - void* h = cumemop_dlopen(); - if (h == nullptr) return false; - // 11.7+ made the v2 entry points the default; older drivers export only the v1 - // names with the same signature. - g_cu_write64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWriteValue64_v2")); - if (g_cu_write64 == nullptr) - g_cu_write64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWriteValue64")); - g_cu_wait64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWaitValue64_v2")); - if (g_cu_wait64 == nullptr) - g_cu_wait64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWaitValue64")); - return g_cu_write64 != nullptr && g_cu_wait64 != nullptr; - }(); - return resolved; -} - -// Functional probe on a scratch pinned int64: enqueue WRITE(7) + WAIT(>=7) + sync. -// Returns true only if the whole memop path works on THIS stream/device/driver. -static bool cumemops_probe(uintptr_t stream, uintptr_t scratch_addr) { - if (!cumemop_resolve()) return false; - auto* s = reinterpret_cast(stream); - if (g_cu_write64(s, (unsigned long long)scratch_addr, 7ULL, kCuWriteDefault) != 0) return false; - if (g_cu_wait64(s, (unsigned long long)scratch_addr, 7ULL, kCuWaitValueGeq) != 0) return false; - return cudaStreamSynchronize(reinterpret_cast(stream)) == cudaSuccess; -} - -// GPU side of the flag handshake (see the block comment above): enqueued on the -// caller's (possibly capturing) stream; the WAIT immediate is the constant 1, -// replay-safe under CUDA graphs. -// The startup probe validates EAGER memops; a driver could still reject them at graph -// capture time. Those enqueue errors would otherwise be swallowed here and surface only -// as a later EndCapture failure -- log the first CUresult so triage is one step. -static void cumemop_check(int rc, const char* what) { - static std::atomic warned{false}; - if (rc != 0 && !warned.exchange(true)) { - std::fprintf(stderr, - "[freetoken/cpu_moe] %s failed with CUresult=%d (first occurrence; " - "subsequent errors are not repeated). If this happened during CUDA " - "graph capture, the driver lacks capture support for stream memops -- " - "set FREETOKEN_CPU_MOE_FLAG_SYNC=0.\n", - what, rc); - } -} - -static void cumemop_submit(uintptr_t stream, uintptr_t done_addr, uintptr_t ready_addr, - int64_t slot) { - auto* s = reinterpret_cast(stream); - // Order matters and is preserved by the front end: reset done BEFORE raising ready, - // so the coordinator's completion write for THIS step can never be wiped. - cumemop_check(g_cu_write64(s, (unsigned long long)(done_addr + (size_t)slot * 8), 0ULL, - kCuWriteDefault), - "cuStreamWriteValue64(done)"); - cumemop_check(g_cu_write64(s, (unsigned long long)(ready_addr + (size_t)slot * 8), 1ULL, - kCuWriteDefault), - "cuStreamWriteValue64(ready)"); -} - -static void cumemop_sync(uintptr_t stream, uintptr_t done_addr, int64_t slot) { - cumemop_check(g_cu_wait64(reinterpret_cast(stream), - (unsigned long long)(done_addr + (size_t)slot * 8), 1ULL, - kCuWaitValueGeq), - "cuStreamWaitValue64(done)"); -} - -struct DotChoice { - dot_fn fn; - const char* name; -}; - -// SIMD tiers, ascending. Each format picks the highest tier <= the one chosen by -// pick_isa() that it implements (fp4 and fp8_block formats have no bf16-specific tier, so the -// avx512bf16 tier maps to their avx512 kernel). -enum IsaTier { ISA_SCALAR = 0, ISA_AVX2 = 1, ISA_AVX512 = 2, ISA_AVX512BF16 = 3 }; - -// Best tier the CPU+build supports, optionally capped DOWN by -// FREETOKEN_CPU_MOE_ISA={scalar,avx2,avx512,avx512bf16} (A/B testing on a machine -// that supports more). FREETOKEN_CPU_MOE_SCALAR=1 forces scalar (legacy alias). -inline IsaTier pick_isa() { -#if CPU_MOE_X86 - if (getenv("FREETOKEN_CPU_MOE_SCALAR")) return ISA_SCALAR; - IsaTier best = ISA_SCALAR; - if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma")) best = ISA_AVX2; - if (best >= ISA_AVX2 && __builtin_cpu_supports("avx512f")) best = ISA_AVX512; -#ifdef CPU_MOE_HAS_AVX512BF16 - if (best >= ISA_AVX512 && __builtin_cpu_supports("avx512bf16")) best = ISA_AVX512BF16; -#endif - if (const char* f = getenv("FREETOKEN_CPU_MOE_ISA")) { - IsaTier want = best; - if (!std::strcmp(f, "scalar")) want = ISA_SCALAR; - else if (!std::strcmp(f, "avx2")) want = ISA_AVX2; - else if (!std::strcmp(f, "avx512")) want = ISA_AVX512; - else if (!std::strcmp(f, "avx512bf16")) want = ISA_AVX512BF16; - if (want < best) best = want; // cap downward; never force above hw/build support - } - return best; -#else - return ISA_SCALAR; -#endif -} - -DotChoice select_dot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 -#ifdef CPU_MOE_HAS_AVX512BF16 - if (t >= ISA_AVX512BF16) return {dot_avx512bf16, "avx512bf16"}; -#endif - if (t >= ISA_AVX512) return {dot_avx512f, "avx512f"}; - if (t >= ISA_AVX2) return {dot_avx2, "avx2"}; -#endif - (void)t; - return {dot_scalar, "scalar"}; -} - -nvdot_fn select_nvdot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (t >= ISA_AVX512) return dot_nvfp4_avx512; - if (t >= ISA_AVX2) return dot_nvfp4_avx2; -#endif - (void)t; - return dot_nvfp4_scalar; -} - -// AVX-VNNI (VEX-256 VPDPBUSD) availability: Alder/Raptor Lake, Sapphire Rapids+, Zen5. -// Distinct from AVX-512 VNNI. Opt out with FREETOKEN_CPU_MOE_NO_VNNI=1 (A/B the W4A8 path). -inline bool cpu_has_avxvnni() { -#if CPU_MOE_X86 - const char* no = getenv("FREETOKEN_CPU_MOE_NO_VNNI"); - if (no && no[0] && no[0] != '0') return false; // ignore unset/empty/"0" - return __builtin_cpu_supports("avx2") && __builtin_cpu_supports("avxvnni"); -#else - return false; -#endif -} +namespace { -// AVX-512 VNNI (512-bit VPDPBUSD): Cascade Lake+, Ice Lake, Sapphire/Emerald Rapids, -// Zen4+. 2x the 256-bit AVX-VNNI width. FREETOKEN_CPU_MOE_NO_AVX512VNNI=1 forces the -// 256-bit path (A/B the two W4A8 kernels on the same box); FREETOKEN_CPU_MOE_NO_VNNI=1 -// still disables the whole W4A8 family (back to the faithful fp32 nvdot), so it is -// honored here too. -inline bool cpu_has_avx512vnni() { -#if CPU_MOE_X86 && defined(CPU_MOE_HAS_AVX512VNNI) - const char* no = getenv("FREETOKEN_CPU_MOE_NO_AVX512VNNI"); - if (no && no[0] && no[0] != '0') return false; - const char* no_vnni = getenv("FREETOKEN_CPU_MOE_NO_VNNI"); - if (no_vnni && no_vnni[0] && no_vnni[0] != '0') return false; - return __builtin_cpu_supports("avx512vnni"); +inline void pin_this_thread(int cpu) { +#if CPU_MOE_HAS_AFFINITY + cpu_set_t set; + CPU_ZERO(&set); + CPU_SET(cpu, &set); + pthread_setaffinity_np(pthread_self(), sizeof(set), &set); #else - return false; -#endif -} - -// Best W4A8 (int8-activation) nvfp4 dot, or nullptr if no SIMD VNNI (caller keeps the -// faithful fp32 nvdot path). The scalar i8 dot exists only as a correctness reference. -nvi8dot_fn select_nvi8dot() { -#if CPU_MOE_X86 -#if defined(CPU_MOE_HAS_AVX512VNNI) - if (cpu_has_avx512vnni()) return dot_nvfp4_i8_avx512vnni; -#endif - if (cpu_has_avxvnni()) return dot_nvfp4_i8_vnni; -#endif - return nullptr; -} - -// ----------------------- DeepSeek-V4 ds_fp4 (W4A8) --------------------------- -// Row-major e2m1 (2/byte, low nibble first) + e8m0 per-32 block scale, no global -// (w = E2M1[code] * 2^(e8m0-127)); activations are FP8-e4m3 round-tripped (per-128 -// block, ue8m0 scale) before each GEMM. Matches kernel/triton/dsv4 (fused_moe + -// fp8_linear): silu(clamp(gate,max=lim)) * clamp(up,-lim,lim), router weight on the -// down output. - -// Activations are pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]) once per -// token/route and reused across all output rows. This drops the hot dot to one -// vpmovzxbd + two vpermps per 32 weights (1.5 shuffle ops / 16 vs 4 for a bf16, -// dup-permute, per-element gather), so the row-major fp4 GEMV stops being port-5 -// bound and approaches the bf16 memory-bandwidth ceiling. -using dsdot_fn = float (*)(const uint8_t*, const uint8_t*, const float*, const float*, int, - const float*, const float*); - -float dot_dsfp4_scalar(const uint8_t* packed, const uint8_t* scale, const float* xe, - const float* xo, int K, const float* e2m1, const float* e8m0) { - float acc = 0.0f; - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const float sc = e8m0[scale[b]]; - const uint8_t* pk = packed + (size_t)b * 16; // 16 bytes = 32 codes - const float* xeb = xe + (size_t)b * 16; - const float* xob = xo + (size_t)b * 16; - float bsum = 0.0f; - for (int j = 0; j < 16; ++j) { - const uint8_t byte = pk[j]; - bsum += e2m1[byte & 0xF] * xeb[j]; // low nibble -> even-K activation - bsum += e2m1[byte >> 4] * xob[j]; // high nibble -> odd-K activation - } - acc += sc * bsum; - } - return acc; -} - -#if CPU_MOE_X86 -// One 32-block: 16 bytes -> 16 low + 16 high nibble values via two vpermps, times -// the pre-split even/odd fp32 activations, folded by the per-32 e8m0 scale. -__attribute__((target("avx512f"))) -inline __m512 dsfp4_blk(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut, - __m512i loma, float sc) { - __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast(pk))); - __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); - __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); - __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob))); - return _mm512_mul_ps(prod, _mm512_set1_ps(sc)); -} - -__attribute__((target("avx512f"))) -float dot_dsfp4_avx512(const uint8_t* packed, const uint8_t* scale, const float* xe, - const float* xo, int K, const float* e2m1, const float* e8m0) { - const __m512 lut = _mm512_loadu_ps(e2m1); - const __m512i loma = _mm512_set1_epi32(0xF); - __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); - const int nb = K / 32; // 16 packed bytes + one e8m0 scale per block - int b = 0; - for (; b + 2 <= nb; b += 2) { // two independent accumulators hide FMA latency - acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16, - xo + (size_t)b * 16, lut, loma, e8m0[scale[b]])); - acc1 = _mm512_add_ps(acc1, dsfp4_blk(packed + (size_t)(b + 1) * 16, xe + (size_t)(b + 1) * 16, - xo + (size_t)(b + 1) * 16, lut, loma, e8m0[scale[b + 1]])); - } - for (; b < nb; ++b) - acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16, - xo + (size_t)b * 16, lut, loma, e8m0[scale[b]])); - return _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1)); -} - -// AVX2: a 32-K block is 16 bytes -> two 8-lane halves (8 even + 8 odd each). -__attribute__((target("avx2,fma"))) -inline __m256 dsfp4_half_avx2(const uint8_t* pk, const float* xeb, const float* xob, __m256 mag8) { - __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); - __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); - __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); - return _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); -} - -__attribute__((target("avx2,fma"))) -float dot_dsfp4_avx2(const uint8_t* packed, const uint8_t* scale, const float* xe, - const float* xo, int K, const float* e2m1, const float* e8m0) { - const __m256 mag8 = _mm256_loadu_ps(e2m1); - __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* pk = packed + (size_t)b * 16; - const float* xeb = xe + (size_t)b * 16; - const float* xob = xo + (size_t)b * 16; - const __m256 sc = _mm256_set1_ps(e8m0[scale[b]]); - acc0 = _mm256_fmadd_ps(dsfp4_half_avx2(pk, xeb, xob, mag8), sc, acc0); - acc1 = _mm256_fmadd_ps(dsfp4_half_avx2(pk + 8, xeb + 8, xob + 8, mag8), sc, acc1); - } - return hsum256(_mm256_add_ps(acc0, acc1)); -} -#endif - -dsdot_fn select_dsdot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (t >= ISA_AVX512) return dot_dsfp4_avx512; - if (t >= ISA_AVX2) return dot_dsfp4_avx2; -#endif - (void)t; - return dot_dsfp4_scalar; -} - -// ------------------------- mxfp4 (gpt-oss) GEMV ----------------------------- -// Transposed split-K layout: blk[Kpairs, N2] (N innermost), scl[Kpairs/16, N2] -// e8m0 per 32-K. Computes out[c] = sum_kb (E2M1[lo]*x[2kb] + E2M1[hi]*x[2kb+1]) -// * 2^(e8m0-127) for a contiguous column tile (blk/scl already offset to col 0 of -// the tile). Vectorized over N (16 columns / __m512), K stays the outer (cache- -// sequential) loop. Used by both gate_up (K=H) and down (K=I). -using mxgemv_fn = void (*)(float*, const uint8_t*, const uint8_t*, const bf16_t*, int, int, - int, const float*, const float*); - -void mxfp4_gemv_scalar(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - for (int c = 0; c < ncol; ++c) out[c] = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t* w = blk + (size_t)kb * N2; - const uint8_t* s = scl + (size_t)(kb >> 4) * N2; - const float xl = bf16_to_f32(x[2 * kb]); - const float xh = bf16_to_f32(x[2 * kb + 1]); - for (int c = 0; c < ncol; ++c) { - const uint8_t byte = w[c]; - out[c] += (e2m1[byte & 0xF] * xl + e2m1[byte >> 4] * xh) * e8m0[s[c]]; - } - } -} - -#if CPU_MOE_X86 -__attribute__((target("avx512f"))) -void mxfp4_gemv_avx512(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - (void)e8m0; // e8m0[c]=2^(c-127) computed via bit construction (no gather) - const __m512 lut = _mm512_loadu_ps(e2m1); - const __m512i loma = _mm512_set1_epi32(0xF); - // K-outer / N-inner: each kb cache line is read once and all live column chunks - // (up to 4 -> 64 cols) accumulate from registers, so DRAM/L2 stream the tile once. - int c0 = 0; - for (; c0 + 16 <= ncol; c0 += 64) { - const int nchunk = std::min(4, (ncol - c0) / 16); - __m512 acc[4]; - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_setzero_ps(); - for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = 32 K = one scale row - __m512 sc[4]; - for (int ci = 0; ci < nchunk; ++ci) { - __m128i sraw = _mm_loadu_si128(reinterpret_cast( - scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 16)); - sc[ci] = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu8_epi32(sraw), 23)); - } - __m512 blk_acc[4]; - for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm512_setzero_ps(); - for (int kk = 0; kk < 16; ++kk) { - const int kb = kblk + kk; - const uint8_t* wbase = blk + (size_t)kb * N2 + c0; - // The transposed layout strides K by N2 bytes; prefetch ahead so the strided - // reads are not exposed to DRAM latency (the HW streamer misses big strides). - constexpr int PFD = 8; - if (kb + PFD < Kpairs) - _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), - _MM_HINT_T0); - const __m512 xl = _mm512_set1_ps(bf16_to_f32(x[2 * kb])); - const __m512 xh = _mm512_set1_ps(bf16_to_f32(x[2 * kb + 1])); - for (int ci = 0; ci < nchunk; ++ci) { - __m512i wi = _mm512_cvtepu8_epi32( - _mm_loadu_si128(reinterpret_cast(wbase + ci * 16))); - __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); - __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); - blk_acc[ci] = _mm512_fmadd_ps(vlo, xl, blk_acc[ci]); - blk_acc[ci] = _mm512_fmadd_ps(vhi, xh, blk_acc[ci]); - } - } - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); - } - for (int ci = 0; ci < nchunk; ++ci) _mm512_storeu_ps(out + c0 + ci * 16, acc[ci]); - } - for (int c = c0; c < ncol; ++c) { // tail columns (< 16) - float o = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t byte = blk[(size_t)kb * N2 + c]; - uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; - float sc; - std::memcpy(&sc, &bits, 4); - o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + - e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; - } - out[c] = o; - } -} - -__attribute__((target("avx2,fma"))) -void mxfp4_gemv_avx2(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - (void)e8m0; // e8m0[s]=2^(s-127) built via s<<23 (no gather) - const __m256 mag8 = _mm256_loadu_ps(e2m1); - int c0 = 0; - for (; c0 + 8 <= ncol; c0 += 32) { // up to 4 chunks of 8 = 32 cols - const int nchunk = std::min(4, (ncol - c0) / 8); - __m256 acc[4]; - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_setzero_ps(); - for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = one scale row - __m256 sc[4]; - for (int ci = 0; ci < nchunk; ++ci) { - __m128i sraw = _mm_loadl_epi64(reinterpret_cast( - scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 8)); - sc[ci] = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu8_epi32(sraw), 23)); - } - __m256 blk_acc[4]; - for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm256_setzero_ps(); - for (int kk = 0; kk < 16; ++kk) { - const int kb = kblk + kk; - const uint8_t* wbase = blk + (size_t)kb * N2 + c0; - constexpr int PFD = 8; - if (kb + PFD < Kpairs) - _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), - _MM_HINT_T0); - const __m256 xl = _mm256_set1_ps(bf16_to_f32(x[2 * kb])); - const __m256 xh = _mm256_set1_ps(bf16_to_f32(x[2 * kb + 1])); - for (int ci = 0; ci < nchunk; ++ci) { - __m256i wi = _mm256_cvtepu8_epi32( - _mm_loadl_epi64(reinterpret_cast(wbase + ci * 8))); - __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); - __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); - blk_acc[ci] = _mm256_fmadd_ps(vlo, xl, blk_acc[ci]); - blk_acc[ci] = _mm256_fmadd_ps(vhi, xh, blk_acc[ci]); - } - } - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); - } - for (int ci = 0; ci < nchunk; ++ci) _mm256_storeu_ps(out + c0 + ci * 8, acc[ci]); - } - for (int c = c0; c < ncol; ++c) { // tail columns (< 8); none when ncol%8==0 - float o = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t byte = blk[(size_t)kb * N2 + c]; - uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; - float sc; - std::memcpy(&sc, &bits, 4); - o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + - e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; - } - out[c] = o; - } -} -#endif - -fp8_block_dot_fn select_fp8_block_dot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - // the avx512bf16 tier keeps the fp32 kernel: on Sapphire Rapids it is faster than widening to bf16 for dpbf16 - if (t >= ISA_AVX512) return dot_fp8_block_avx512f; - if (t >= ISA_AVX2) return dot_fp8_block_avx2; + (void)cpu; #endif - (void)t; - return dot_fp8_block_scalar; } -mxgemv_fn select_mxgemv() { - const IsaTier t = pick_isa(); +inline void cpu_relax() { #if CPU_MOE_X86 - if (t >= ISA_AVX512) return mxfp4_gemv_avx512; - if (t >= ISA_AVX2) return mxfp4_gemv_avx2; + _mm_pause(); +#elif defined(__aarch64__) + asm volatile("yield" ::: "memory"); #endif - (void)t; - return mxfp4_gemv_scalar; -} - -// Round a clamped |x|<=448 to nearest float8-e4m3 (RNE), back to fp32. Matches -// torch.float8_e4m3fn / triton .to(float8e4nv). -inline float e4m3_round(float x) { - const float sign = x < 0.0f ? -1.0f : 1.0f; - const float a = std::fabs(x); - if (a == 0.0f) return 0.0f; - if (a >= 448.0f) return sign * 448.0f; - int e; - std::frexp(a, &e); // a in [2^(e-1), 2^e) - float step = std::ldexp(1.0f, e - 4); - const float min_step = std::ldexp(1.0f, -9); // e4m3 subnormal step (2^-9) - if (step < min_step) step = min_step; - float r = std::nearbyint(a / step) * step; - if (r > 448.0f) r = 448.0f; - return sign * r; -} - -// IEEE ceil(log2(v)) for v>0 (matches dsv4 _log2_ceil / fast_round_scale). -inline int ceil_log2_pos(float v) { - uint32_t bits; - std::memcpy(&bits, &v, sizeof(bits)); - const int exp = (int)((bits >> 23) & 0xFF); - const int man = (int)(bits & 0x7FFFFF); - return exp - 127 + (man != 0 ? 1 : 0); -} - -// Split an interleaved bf16 row into fp32 even/odd halves (even[m]=src[2m]). -// bf16->fp32 is exact, so this only reorders -- done once per token/route and -// reused across every output row of the GEMV. -inline void deinterleave_bf16_f32(const bf16_t* src, float* even, float* odd, int K) { - for (int m = 0; m < K / 2; ++m) { - even[m] = bf16_to_f32(src[2 * m]); - odd[m] = bf16_to_f32(src[2 * m + 1]); - } -} - -// DeepSeek-V4 / V4.1 activation FP8 round-trip (bf16 in/out): per `block` (the checkpoint's -// fp8 block: 128 on V4, 32 on V4.1), s = 2^ceil(log2(max(|x|,1e-4)/448)); -// y = round_e4m3(clamp(x/s,+-448)) * s. -void fp8_roundtrip_bf16(const bf16_t* src, bf16_t* dst, int K, int block) { - for (int b0 = 0; b0 < K; b0 += block) { - const int b1 = std::min(K, b0 + block); - float amax = 1e-4f; - for (int i = b0; i < b1; ++i) amax = std::max(amax, std::fabs(bf16_to_f32(src[i]))); - const float s = std::ldexp(1.0f, ceil_log2_pos(amax * (1.0f / 448.0f))); - const float inv_s = 1.0f / s; - for (int i = b0; i < b1; ++i) { - float q = bf16_to_f32(src[i]) * inv_s; - q = std::min(448.0f, std::max(-448.0f, q)); - dst[i] = f32_to_bf16(e4m3_round(q) * s); - } - } } -// --------------------------------- executor --------------------------------- - -struct CpuMoeExecutor; +constexpr int kMaxTopK = 64; +// One MoE layer on the first num_tokens rows of the pinned IO. +// It is also the doorbell word the GPU writes, a constant baked into each captured graph; num_tokens >= 1 keeps it nonzero, since 0 means no doorbell. struct MoeTask { - CpuMoeExecutor* exec; - int layer_id; - int num_tokens; - const bf16_t* x; // [num_tokens, H] - const int32_t* ids; // [num_tokens, top_k] (raw expert ids; <0 = skip) - const float* w; // [num_tokens, top_k] - bf16_t* y; // [num_tokens, H] + uint64_t layer : 32; + uint64_t num_tokens : 32; }; +static_assert(sizeof(MoeTask) == sizeof(uint64_t), "MoeTask must pack into one 64-bit word"); -// Output-row tiling. Small enough to give every worker independent work even at -// batch size 1; large enough to amortize the atomic work-grab. -// -// Bandwidth notes (Sapphire Rapids 8480+, 13 cores): the two passes already read -// every expert weight byte exactly once per token (each output row block is owned -// by one worker), and x stays hot in L1 across a (token,expert)'s rows -- so the -// kernel is single-read bandwidth-optimal at bs=1 (~205 GB/s vs ~55 GB/s PCIe). -// One worker per *physical* core, pinned, is the sweet spot; SMT oversubscription -// thrashes the spin-barrier. Deferred (not worth it here / for this workload): -// - AMX-bf16: a GEMM tile engine; decode is M=1 GEMV so tiles sit idle. It would -// only pay off in a grouped/batched (dedup) path. -// - expert dedup for bs>1: read each distinct expert once and GEMM its tokens. -// Helps locality+bytes when bs is large; decode batches here are tiny (<=4). -// - NUMA: a single node is assumed. Multi-socket machines would split each -// expert's K dimension per node (banks are already per-row contiguous). -constexpr int IBLK = 32; -constexpr int HBLK = 32; - -// -------------------------------- Q4_0 (W4A8) -------------------------------- -// Native GGUF Q4_0 experts (gemma4 GGUF): per-32 block = fp16 scale d + 16 packed -// bytes; byte j holds element j in its low nibble and j+16 in its high nibble, so a -// block's storage order is [lo0..lo15, hi0..hi15] and w = (nibble - 8) * d. Matches -// the reference dequant (models/gguf/dequant.py) and the packed banks the GPU offload -// path streams. -// -// llama.cpp ggml_vec_dot_q4_0_q8_0: W4A8. The activation is pre-quantized to Q8_0 -// (per-32-block int8 ``aq`` + fp32 scale ``asb``); each block unpacks its 16 bytes to -// 32 int8 weights in [-8,7] (bytes_from_nibbles_32: low nibbles -> elems 0..15, high -// -> 16..31) and runs an integer block dot -- VPDPBUSD (AVX-VNNI) or VPMADDUBSW+VPMADDWD -// (AVX2) with the ggml sign trick |w|*(sign(w)*a)=w*a, or a scalar int loop -- then -// scales the block sum by wd*xd in fp32. No fp weight dequant / shuffle chain. The GPU -// offload path (ggml_moe_a8_vec / MMVQ) is also W4A8, so cpu and hybrid stay close. -using q4dot_fn = float (*)(const uint8_t*, const int8_t*, const float*, int); - -float q4_0_dot_i8_scalar(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - float acc = 0.0f; - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - const uint8_t* q = blk + 2; // 16 nibble bytes - const int8_t* a = aq + (size_t)b * 32; - int isum = 0; - for (int j = 0; j < 16; ++j) { - isum += ((int)(q[j] & 0x0F) - 8) * (int)a[j]; // elem j - isum += ((int)(q[j] >> 4) - 8) * (int)a[16 + j]; // elem 16+j - } - acc += fp16_to_f32(dh) * asb[b] * (float)isum; - } - return acc; -} - -#if CPU_MOE_X86 -// fp16 block scale -> fp32 via HW F16C (single value in lane 0). -__attribute__((target("f16c"))) -static inline float q4_scale(uint16_t h) { - return _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128((int)h))); +inline uint64_t to_word(MoeTask t) { + uint64_t w; + std::memcpy(&w, &t, sizeof(w)); + return w; } -// Unpack one Q4_0 block's 16 bytes -> 32 int8 weights in [-8,7] (elems 0..15 = low -// nibbles, 16..31 = high nibbles). ``eight`` = _mm256_set1_epi8(8). -__attribute__((target("avx2"))) -static inline __m256i q4_unpack32(const uint8_t* blk, __m128i mask, __m256i eight) { - const __m128i qb = _mm_loadu_si128(reinterpret_cast(blk + 2)); - const __m128i lo = _mm_and_si128(qb, mask); - const __m128i hi = _mm_and_si128(_mm_srli_epi16(qb, 4), mask); - return _mm256_sub_epi8(_mm256_set_m128i(hi, lo), eight); +inline MoeTask from_word(uint64_t w) { + MoeTask t; + std::memcpy(&t, &w, sizeof(t)); + return t; } -// AVX2 W4A8 (llama.cpp non-VNNI mul_sum_i8_pairs): integer block dot via VPMADDUBSW + -// VPMADDWD (sign trick), scaled by wd*xd. |aw*sa| pair sums <= 8*127*2 < 32767 -> no -// int16 saturation. This is the fast path on AVX2 CPUs without AVX-VNNI (and the -// avx512-tier fallback, since the block dot is 256-bit either way). -__attribute__((target("avx2,fma,f16c"))) -float q4_0_dot_i8_avx2(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - const __m128i mask = _mm_set1_epi8(0x0F); - const __m256i eight = _mm256_set1_epi8(8); - const __m256i ones16 = _mm256_set1_epi16(1); - __m256 accF = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - __m256i wq = q4_unpack32(blk, mask, eight); - __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); - __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) - __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) - __m256i d32 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, sa), ones16); // 8 int32 - accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d32), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); - } - return hsum256(accF); -} - -// AVX-VNNI W4A8: one VPDPBUSD per block (the fast path on modern CPUs). -__attribute__((target("avx2,avxvnni,fma,f16c"))) -float q4_0_dot_i8_vnni(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - const __m128i mask = _mm_set1_epi8(0x0F); - const __m256i eight = _mm256_set1_epi8(8); - __m256 accF = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - __m256i wq = q4_unpack32(blk, mask, eight); - __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); - __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) - __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) - __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); - // All 32 elems of the block share wd*xd; distribute over di's 8 partial sums and - // reduce at the end (equivalent to scale * block_total). - accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); - } - return hsum256(accF); -} -#endif // CPU_MOE_X86 - -// All tiers are W4A8 (int8 activations pre-quantized to Q8_0). AVX-VNNI is orthogonal to -// the ISA tier (gated by cpu_has_avxvnni() / FREETOKEN_CPU_MOE_NO_VNNI), so it wins when -// present; otherwise the 256-bit VPMADDUBSW kernel covers both the avx2 and avx512 tiers. -q4dot_fn select_q4dot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (cpu_has_avxvnni()) return q4_0_dot_i8_vnni; - if (t >= ISA_AVX2) return q4_0_dot_i8_avx2; -#endif - (void)t; - return q4_0_dot_i8_scalar; -} - -enum WFmt { WF_BF16 = 0, WF_NVFP4 = 1, WF_MXFP4 = 2, WF_DSFP4 = 3, WF_Q4_0 = 4, - WF_FP8_BLOCK = 5 }; +// Every weight format this module serves, keyed by the name cpu_executor.py passes; weight_formats() exports the names. +const std::pair (*)(const MoeParams&, const FormatArgs&)> kFormats[] = { + {"bf16", make_bf16_format}, {"nvfp4", make_nvfp4_format}, {"mxfp4_triton", make_mxfp4_format}, + {"ds_fp4", make_ds_fp4_format}, {"q4_0", make_q4_0_format}, {"fp8_block", make_fp8_block_format}, +}; -// Each ctor pointer arg is the address of a CPU int64 array of length -// num_layers (one base address per layer, built by cpu_executor.py's -// _make_table), not a single flat bank. tbl_at resolves -// tbl[layer_id] once per task/pass; a null table (bank unused by this fmt, ptr -// arg 0) resolves to nullptr without dereferencing. -inline const void* tbl_at(const uint64_t* tbl, int layer_id) { - return tbl ? reinterpret_cast(tbl[layer_id]) : nullptr; +std::unique_ptr make_format(const std::string& name, const MoeParams& p, const FormatArgs& a) { + for (const auto& [n, make] : kFormats) + if (name == n) return make(p, a); + throw std::invalid_argument("unknown CPU MoE weight format " + name); } struct CpuMoeExecutor { + // A host-function callback gets one pointer, so it carries the executor with the task. + struct HostTask { + CpuMoeExecutor* exec; + MoeTask task; + }; + + const int num_layers, num_experts; + const bool apply_on_input; + MoeParams params; // the format keeps a reference; declared before fmt + std::unique_ptr fmt; + bool preps_inter = false; + bool rounds_routes = false; + int n_iblk = 0, n_hblk = 0; // kRowBlock blocks per intermediate row and per output row + + // Pinned IO shared by every layer and batch size (one task in flight): [max_tokens, H] x and y, [max_tokens, top_k] ids and w. + const bf16_t* io_x; + const int32_t* io_ids; + const float* io_w; + bf16_t* io_y; + AlignedVector g_scratch; // [max_tokens * top_k * I] intermediate + + // Worker pool. int num_threads; - int num_layers, num_experts, top_k; - int H, I; - int act, apply_on_input; - int fmt; // WFmt - bool needs_di = false; // pre-deinterleave activations to fp32 (nvfp4/ds_fp4/fp8_block) - // Per-layer pointer tables (one base address per layer, see tbl_at). gate_up_tbl - // doubles as the bf16 gate_up table and the nvfp4/mxfp4/q4_0/ds_fp4 packed-gate_up - // table (down_tbl likewise for down); which reinterpretation applies is picked by - // fmt at each resolve site (see gemm1_dot/gemm2_dot/do_pass1_mxfp4/do_pass1_dsfp4). - const uint64_t* gate_up_tbl; // bf16: [E,2I,H] rows; else: packed e2m1/mxfp4-blocks - const uint64_t* down_tbl; // bf16: [E,H,I] rows; else: packed e2m1/mxfp4-blocks - const uint64_t* gu_scale_tbl; // nvfp4/mxfp4/ds_fp4: [E,2I,*] block scales - const uint64_t* gu_global_tbl; // nvfp4: [E,2I] fp16 row globals - const uint64_t* dn_scale_tbl; // nvfp4/mxfp4/ds_fp4: [E,H,*] block scales - const uint64_t* dn_global_tbl; // nvfp4: [E,H] fp16 row globals - const uint64_t* gu_bias_tbl; // mxfp4: [E,2I] bf16 biases - const uint64_t* dn_bias_tbl; // mxfp4: [E,H] bf16 biases - float swiglu_alpha; - float swiglu_limit; // +inf == no clamp - dot_fn dot; - nvdot_fn nvdot; - nvi8dot_fn nvi8dot = nullptr; // AVX-VNNI W4A8 nvfp4 dot (nullptr -> use fp32 nvdot) - bool use_vnni = false; // nvfp4 + AVX-VNNI: decode via int8 VPDPBUSD (W4A8) - bool use_q4a8 = false; // q4_0: always W4A8 (llama.cpp Q4_0 x Q8_0); int8 pre-quant - dsdot_fn dsdot; - mxgemv_fn mxgemv; - q4dot_fn q4dot; - fp8_block_dot_fn fp8_block_dot; - // fp8_block: bf16 scale per 128x128 weight block. These are the scale tensor's - // trailing strides, including any padding used by the host bank layout. - int fp8_block_gu_scale_row = 0, fp8_block_dn_scale_row = 0; - // ds_fp4: the caller already FP8-round-tripped the input activations on the GPU - // (same reference grid), so submit() must not repeat it on the host-callback - // thread. That scalar per-element pass is single-threaded ON THE DECODE CRITICAL - // PATH (~0.3ms/layer at H=4096, every worker and the GPU waiting on it); moving - // it to a captured GPU elementwise kernel removes it while keeping the official - // W4A8 numerics bit-exact. Set via set_input_prequant (see cpu_executor.py). - bool input_prequant = false; - // ds_fp4 activation round-trip block (the checkpoint's fp8 block); set via set_act_block. - int act_block = 128; - // Q4_0 packed-row byte strides (H/32*18 for gate_up over K=H, I/32*18 for down over K=I). - int q4_gu_row_bytes = 0, q4_dn_row_bytes = 0; - float e2m1_lut[16]; - float e4m3_lut[256]; - float e8m0_lut[256]; // mxfp4 block scale: 2^(s-127), s clamped to [0,254] - const char* isa; - - std::vector g_scratch; // [max_tokens * top_k * I] intermediate - std::vector xq_scratch; // [max_tokens * H] ds_fp4 fp8-roundtripped input - // nvfp4/ds_fp4/fp8_block activations pre-deinterleaved to fp32 (even/odd K) for the row-major dot. - std::vector xe_scratch, xo_scratch; // [max_tokens * H/2] (input) - std::vector ge_scratch, go_scratch; // [max_tokens*top_k*I/2] (intermediate) - // AVX-VNNI W4A8: per-16-block int8 activations [even(8),odd(8)] + per-block scale. - std::vector xi8_scratch, gi8_scratch; // [max_tokens*H], [max_tokens*top_k*I] - std::vector xas_scratch, gas_scratch; // [max_tokens*H/16], [..*top_k*I/16] - std::string isa_str; - + std::vector core_ids; // worker tid -> logical CPU to pin to (may be empty) std::vector workers; std::mutex task_mtx; std::condition_variable task_cv; @@ -1421,7 +135,7 @@ struct CpuMoeExecutor { std::atomic stop{false}; // spinning workers read it without the lock uint64_t cur_gen = 0; - MoeTask* cur_task = nullptr; + MoeTask cur_task{}; // Separate cache lines: idle workers poll `submitted` while busy workers hit p1_next/p2_next. alignas(64) std::atomic submitted{0}; alignas(64) std::atomic completed{0}; @@ -1431,264 +145,55 @@ struct CpuMoeExecutor { std::atomic spin_wait{false}; bool coord_pinned = false; + // Work split of the current task. alignas(64) std::atomic p1_next{0}; std::atomic p2_next{0}; - std::atomic prt_next{0}; // intermediate-row prep phase (prep_g_row) + std::atomic prt_next{0}; // intermediate-row prep phase (prep_inter) int64_t p1_total = 0, p2_total = 0, prt_total = 0; - int n_iblk = 0, n_hblk = 0; std::atomic done_count{0}; std::atomic bar_count{0}; std::atomic bar_sense{0}; - std::vector owned_tasks; // persistent task descriptors (graph-stable) - std::vector core_ids; // worker tid -> logical CPU to pin to (may be empty) - // ---- Flag-based GPU<->CPU handshake (replaces the per-layer cudaLaunchHostFunc pair) ---- - // A tiny GPU kernel bumps ready_flags[slot] at submit; this coordinator thread busy-polls - // it, runs the slot's task on the worker pool, and sets done_flags[slot], which a GPU - // spin-wait kernel polls at sync. This removes the ~2x30-50us host-func dispatch round - // trips per MoE layer per decode step that otherwise idle the GPU (~6 ms/step on a - // 75-layer model). One slot per (layer, decode batch size) pair -- the Python side - // allocates slots as tasks are created. Flags live in mapped-pinned host memory (UVA: - // the same pointers are used by the GPU kernels and by this thread). + // The GPU rings the doorbell with a task word at submit; this coordinator thread busy-polls it, runs the task on the worker pool, and completes it, which the GPU waits on at sync. + // This removes the ~2x30-50us host-func dispatch round trips per MoE layer per decode step that otherwise idle the GPU (~6 ms/step on a 75-layer model). + std::unique_ptr gpu_sync; std::thread coord_thread; std::atomic coord_stop{false}; - volatile int64_t* ready_flags = nullptr; // GPU increments, this thread polls - volatile int64_t* done_flags = nullptr; // this thread sets, GPU spin-waits - int coord_num_slots = 0; - std::vector flag_task; // slot -> task (registered lazily) - std::vector flag_served; // slot -> completed dispatch count (tests/debug) - std::mutex flag_task_mtx; - - // Portable ordering for the flag handshake: "ready observed => the DMA'd inputs that - // preceded the bump are visible" and "y stores are visible before done". Plain - // volatile loads lean on x86 TSO; acquire/release makes it hold on aarch64 too - // (GH200/Jetson) at zero x86 cost. (MSVC branch is x86-only today: compiler barrier - // + TSO.) - static int64_t flag_load_acquire(const volatile int64_t* p) { -#if defined(_MSC_VER) - const int64_t v = *p; - _ReadWriteBarrier(); - return v; -#else - return __atomic_load_n(const_cast(p), __ATOMIC_ACQUIRE); -#endif - } - - static void flag_store_release(volatile int64_t* p, int64_t v) { -#if defined(_MSC_VER) - _ReadWriteBarrier(); - *p = v; -#else - __atomic_store_n(const_cast(p), v, __ATOMIC_RELEASE); -#endif - } + std::atomic flag_served{0}; // completed doorbells (tests, watchdog) + std::map host_tasks; // host-function userData, keyed by doorbell word (graph-stable) CpuMoeExecutor(int num_threads_, int num_layers_, int num_experts_, int top_k_, - int hidden_size, int inter_size, int max_tokens, int activation_id, - int apply_router_weight_on_input, int weight_format, - uintptr_t gate_up_ptr, uintptr_t down_ptr, uintptr_t gate_up_scale_ptr, - uintptr_t gate_up_global_ptr, uintptr_t down_scale_ptr, - uintptr_t down_global_ptr, uintptr_t gate_up_bias_ptr, - uintptr_t down_bias_ptr, double swiglu_alpha_, double swiglu_limit_, - std::vector core_ids_, int fp8_block_gu_scale_stride_, - int fp8_block_dn_scale_stride_) - : num_threads(num_threads_ > 0 ? num_threads_ : 1), - num_layers(num_layers_), + int hidden_size, int inter_size, int max_tokens_, int activation_id, + int apply_router_weight_on_input, const std::string& weight_format, + std::map banks, std::map options, + double swiglu_alpha_, double swiglu_limit_, + std::vector core_ids_, uintptr_t x_ptr, uintptr_t ids_ptr, uintptr_t w_ptr, + uintptr_t y_ptr) + : num_layers(num_layers_), num_experts(num_experts_), - top_k(top_k_), - H(hidden_size), - I(inter_size), - act(activation_id), - apply_on_input(apply_router_weight_on_input), - fmt(weight_format), - gate_up_tbl(reinterpret_cast(gate_up_ptr)), - down_tbl(reinterpret_cast(down_ptr)), - gu_scale_tbl(reinterpret_cast(gate_up_scale_ptr)), - gu_global_tbl(reinterpret_cast(gate_up_global_ptr)), - dn_scale_tbl(reinterpret_cast(down_scale_ptr)), - dn_global_tbl(reinterpret_cast(down_global_ptr)), - gu_bias_tbl(reinterpret_cast(gate_up_bias_ptr)), - dn_bias_tbl(reinterpret_cast(down_bias_ptr)), - swiglu_alpha(static_cast(swiglu_alpha_)), - swiglu_limit(static_cast(swiglu_limit_)), - fp8_block_gu_scale_row(fp8_block_gu_scale_stride_), - fp8_block_dn_scale_row(fp8_block_dn_scale_stride_), + apply_on_input(apply_router_weight_on_input != 0), + params{top_k_, max_tokens_, hidden_size, inter_size, activation_id, static_cast(swiglu_alpha_), + static_cast(swiglu_limit_)}, + io_x(reinterpret_cast(x_ptr)), + io_ids(reinterpret_cast(ids_ptr)), + io_w(reinterpret_cast(w_ptr)), + io_y(reinterpret_cast(y_ptr)), + num_threads(num_threads_ > 0 ? num_threads_ : 1), core_ids(std::move(core_ids_)) { - DotChoice c = select_dot(); - dot = c.fn; - nvdot = select_nvdot(); - dsdot = select_dsdot(); - mxgemv = select_mxgemv(); - q4dot = select_q4dot(); - fp8_block_dot = select_fp8_block_dot(); - if (weight_format == WF_Q4_0) { - if (H % 32 != 0 || I % 32 != 0) - throw std::runtime_error("Q4_0 CPU MoE requires H and I to be multiples of 32"); - q4_gu_row_bytes = (H / 32) * 18; // K = H (gate_up rows) - q4_dn_row_bytes = (I / 32) * 18; // K = I (down rows) - } - isa = c.name; - // nvfp4 (AVX-VNNI only): W4A8 int8 decode when the CPU supports it. q4_0 is always - // 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(); - 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") : ""; - const char* vnni_tag = - cpu_has_avx512vnni() ? "+avx512vnni(nvfp4-w4a8)" : "+vnni(nvfp4-w4a8)"; - isa_str = std::string(c.name) + (use_vnni ? vnni_tag : "") + q4tag; - isa = isa_str.c_str(); - for (int i = 0; i < 16; ++i) e2m1_lut[i] = kE2M1[i]; - for (int i = 0; i < 256; ++i) e4m3_lut[i] = e4m3_decode((uint8_t)i); - // e8m0 (mxfp4 block scale) = 2^(s-127); the GPU GEMV clamps s to [0,254]. - for (int i = 0; i < 256; ++i) e8m0_lut[i] = std::ldexp(1.0f, std::min(i, 254) - 127); - g_scratch.assign(static_cast(max_tokens) * top_k * I, 0); - // Row-major fp4 (nvfp4/ds_fp4) and fp8_block pre-deinterleave activations to fp32 even/odd. - needs_di = (fmt == WF_NVFP4 || fmt == WF_DSFP4 || fmt == WF_FP8_BLOCK); - if (needs_di) { - if (fmt == WF_DSFP4) xq_scratch.assign(static_cast(max_tokens) * H, 0); - xe_scratch.assign(static_cast(max_tokens) * (H / 2), 0); - xo_scratch.assign(static_cast(max_tokens) * (H / 2), 0); - ge_scratch.assign(static_cast(max_tokens) * top_k * (I / 2), 0); - go_scratch.assign(static_cast(max_tokens) * top_k * (I / 2), 0); - if (use_vnni) { - xi8_scratch.assign(static_cast(max_tokens) * H, 0); - xas_scratch.assign(static_cast(max_tokens) * (H / 16), 0); - gi8_scratch.assign(static_cast(max_tokens) * top_k * I, 0); - gas_scratch.assign(static_cast(max_tokens) * top_k * (I / 16), 0); - } - } - // q4_0 W4A8: per-32-block Q8_0 activations (int8 + fp32 scale) for input + intermediate. - if (use_q4a8) { - xi8_scratch.assign(static_cast(max_tokens) * H, 0); - xas_scratch.assign(static_cast(max_tokens) * (H / 32), 0); - gi8_scratch.assign(static_cast(max_tokens) * top_k * I, 0); - gas_scratch.assign(static_cast(max_tokens) * top_k * (I / 32), 0); - } + if (params.top_k > kMaxTopK) throw std::invalid_argument("CPU MoE supports top_k <= " + std::to_string(kMaxTopK)); + fmt = make_format(weight_format, params, FormatArgs{weight_format, std::move(banks), std::move(options)}); + preps_inter = fmt->preps_inter(); + rounds_routes = fmt->rounds_routes(); + n_iblk = (params.I + kRowBlock - 1) / kRowBlock; + n_hblk = (params.H + kRowBlock - 1) / kRowBlock; + g_scratch.assign(static_cast(params.max_tokens) * params.top_k * params.I, 0); for (int t = 0; t < num_threads; ++t) workers.emplace_back([this, t] { worker_loop(t); }); } - // Quantize a bf16 activation row to Q8_0 (llama.cpp): per-32-block symmetric int8 in - // natural order + one fp32 scale (amax/127) per block. Done once per token/route, - // amortized over every output row that the q4_0 W4A8 GEMV reads. K % 32 == 0. - void quant_q8_0(const bf16_t* x, int K, int8_t* aq, float* asb) { - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const bf16_t* xb = x + (size_t)b * 32; - float xf[32], amax = 0.0f; - for (int j = 0; j < 32; ++j) { - xf[j] = bf16_to_f32(xb[j]); - amax = std::max(amax, std::fabs(xf[j])); - } - const float d = amax > 0.0f ? amax / 127.0f : 1.0f; - asb[b] = d; - const float inv = amax > 0.0f ? 1.0f / d : 0.0f; - int8_t* o = aq + (size_t)b * 32; - for (int j = 0; j < 32; ++j) - o[j] = (int8_t)std::max(-127, std::min(127, (int)std::lround(xf[j] * inv))); - } - } - - // Quantize the pre-deinterleaved fp32 even/odd activations to per-16-block int8 in the - // [even(8),odd(8)] layout the VNNI dot expects. Done once per token/route (amortized over - // every output row), so a scalar pass is fine relative to the GEMV. - void quant_i8_pg16(const float* xe, const float* xo, int K, int8_t* asi8, float* asb) { - const int nb = K / 16; - for (int b = 0; b < nb; ++b) { - const float* xeb = xe + (size_t)b * 8; - const float* xob = xo + (size_t)b * 8; - float amax = 0.0f; - for (int j = 0; j < 8; ++j) - amax = std::max(amax, std::max(std::fabs(xeb[j]), std::fabs(xob[j]))); - const float s = amax > 0.0f ? amax / 127.0f : 1.0f; - asb[b] = s; - const float inv = 1.0f / s; - int8_t* ae = asi8 + (size_t)b * 16; - for (int j = 0; j < 8; ++j) { - int qe = (int)std::lround(xeb[j] * inv), qo = (int)std::lround(xob[j] * inv); - ae[j] = (int8_t)std::max(-127, std::min(127, qe)); - ae[8 + j] = (int8_t)std::max(-127, std::min(127, qo)); - } - } - } - - // gate_up output row `row` (in [0, 2I)) dotted with activation over K = H. ``e`` is - // the layer-local expert row (0..num_experts); the layer bases (already resolved - // once per task/pass by the caller via tbl_at) pick the layer's own tensors. - // bf16 uses the interleaved bf16 row; nvfp4 uses the pre-split fp32 even/odd halves - // (or, with AVX-VNNI, the per-16-block int8 activations). - // fp8_block uses the halves for whole 128-blocks and the bf16 row for a ragged tail. - inline float gemm1_dot(const bf16_t* gate_up_l, const uint8_t* gu_packed_l, - const uint8_t* gu_scale_l, const uint16_t* gu_global_l, int e, int row, - const bf16_t* x, const float* xe, const float* xo, const int8_t* xi8, - const float* xas) { - if (fmt == WF_BF16) { - const bf16_t* w = gate_up_l + ((size_t)e * (2 * I) + row) * H; - return dot(w, x, H); - } - if (fmt == WF_Q4_0) { - const uint8_t* w = - gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)q4_gu_row_bytes; - return q4dot(w, xi8, xas, H); // W4A8: int8 activations (Q8_0), scale in xas - } - if (fmt == WF_FP8_BLOCK) { - const uint8_t* w = gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)H; - const bf16_t* sc = reinterpret_cast(gu_scale_l) + - ((size_t)e * ((2 * I + FP8_BLK - 1) / FP8_BLK) + row / FP8_BLK) * - (size_t)fp8_block_gu_scale_row; - return fp8_block_dot(w, xe, xo, x, sc, H, e4m3_lut); - } - const size_t r = (size_t)e * (2 * I) + row; - if (use_vnni) - return nvi8dot(gu_packed_l + r * (size_t)(H / 2), gu_scale_l + r * (size_t)(H / 16), - fp16_to_f32(gu_global_l[r]), xi8, H, e4m3_lut, xas); - return nvdot(gu_packed_l + r * (size_t)(H / 2), gu_scale_l + r * (size_t)(H / 16), - fp16_to_f32(gu_global_l[r]), xe, xo, H, e2m1_lut, e4m3_lut); - } - - // down output row `row` (in [0, H)) dotted with the intermediate over K = I. Same - // layer-local-base convention as gemm1_dot. - inline float gemm2_dot(const bf16_t* down_l, const uint8_t* dn_packed_l, - const uint8_t* dn_scale_l, const uint16_t* dn_global_l, int e, int row, - const bf16_t* g, const float* ge, const float* go, const int8_t* gi8, - const float* gas) { - if (fmt == WF_BF16) { - const bf16_t* w = down_l + ((size_t)e * H + row) * I; - return dot(w, g, I); - } - if (fmt == WF_Q4_0) { - const uint8_t* w = dn_packed_l + ((size_t)e * H + row) * (size_t)q4_dn_row_bytes; - return q4dot(w, gi8, gas, I); // W4A8: int8 activations (Q8_0), scale in gas - } - if (fmt == WF_FP8_BLOCK) { - const uint8_t* w = dn_packed_l + ((size_t)e * H + row) * (size_t)I; - const bf16_t* sc = reinterpret_cast(dn_scale_l) + - ((size_t)e * ((H + FP8_BLK - 1) / FP8_BLK) + row / FP8_BLK) * - (size_t)fp8_block_dn_scale_row; - return fp8_block_dot(w, ge, go, g, sc, I, e4m3_lut); - } - const size_t r = (size_t)e * H + row; - if (use_vnni) - return nvi8dot(dn_packed_l + r * (size_t)(I / 2), dn_scale_l + r * (size_t)(I / 16), - fp16_to_f32(dn_global_l[r]), gi8, I, e4m3_lut, gas); - return nvdot(dn_packed_l + r * (size_t)(I / 2), dn_scale_l + r * (size_t)(I / 16), - fp16_to_f32(dn_global_l[r]), ge, go, I, e2m1_lut, e4m3_lut); - } - void pin_self(int tid) { -#if CPU_MOE_HAS_AFFINITY - if (core_ids.empty()) return; - const int cpu = core_ids[tid % static_cast(core_ids.size())]; - cpu_set_t set; - CPU_ZERO(&set); - CPU_SET(cpu, &set); - pthread_setaffinity_np(pthread_self(), sizeof(set), &set); -#else - (void)tid; -#endif + if (!core_ids.empty()) pin_this_thread(core_ids[tid % static_cast(core_ids.size())]); } ~CpuMoeExecutor() { @@ -1701,31 +206,16 @@ struct CpuMoeExecutor { task_cv.notify_all(); for (auto& th : workers) if (th.joinable()) th.join(); - for (MoeTask* t : owned_tasks) delete t; } - uintptr_t create_task(int layer_id, int num_tokens, uintptr_t x_ptr, - uintptr_t ids_ptr, uintptr_t w_ptr, uintptr_t y_ptr) { - MoeTask* t = new MoeTask{this, - layer_id, - num_tokens, - reinterpret_cast(x_ptr), - reinterpret_cast(ids_ptr), - reinterpret_cast(w_ptr), - reinterpret_cast(y_ptr)}; - owned_tasks.push_back(t); - return reinterpret_cast(t); + MoeTask task(int layer_id, int num_tokens) const { + if (layer_id < 0 || layer_id >= num_layers || num_tokens < 1 || num_tokens > params.max_tokens) + throw std::out_of_range("CPU MoE task (layer " + std::to_string(layer_id) + ", " + + std::to_string(num_tokens) + " tokens) is outside the executor's layers or IO"); + return MoeTask{static_cast(layer_id), static_cast(num_tokens)}; } - const char* isa_name() const { return isa; } - - const char* fp8_block_isa_name() const { -#if CPU_MOE_X86 - if (fp8_block_dot == dot_fp8_block_avx512f) return "avx512f"; - if (fp8_block_dot == dot_fp8_block_avx2) return "avx2"; -#endif - return "scalar"; - } + std::string isa_name() const { return fmt->isa_name(); } void barrier(int& local_sense) { local_sense ^= 1; @@ -1733,275 +223,58 @@ struct CpuMoeExecutor { bar_count.store(0); bar_sense.store(local_sense); } else { - while (bar_sense.load() != local_sense) { -#if CPU_MOE_X86 - _mm_pause(); -#endif - } + while (bar_sense.load() != local_sense) cpu_relax(); } } - void do_pass1(const MoeTask* t, int64_t p) { - if (fmt == WF_MXFP4) { - do_pass1_mxfp4(t, p); - return; - } - if (fmt == WF_DSFP4) { - do_pass1_dsfp4(t, p); - return; - } + // Pass 1 item p = one kRowBlock block of one (token, route)'s intermediate row. + void do_pass1(const MoeTask& t, int64_t p) { + const int top_k = params.top_k, H = params.H, I = params.I; const int64_t ib = p % n_iblk; const int64_t tk = p / n_iblk; const int k = static_cast(tk % top_k); const int tok = static_cast(tk / top_k); - const int e = t->ids[static_cast(tok) * top_k + k]; + const int e = io_ids[static_cast(tok) * top_k + k]; if (e < 0 || e >= num_experts) return; - const float w_in = apply_on_input ? t->w[static_cast(tok) * top_k + k] : 1.0f; - // Resolve this task's layer base once; row indexing below is layer-local (e). - 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 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; - bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I; - const int i0 = static_cast(ib) * IBLK; - const int i1 = std::min(I, i0 + IBLK); - const bool clamped = act == ACT_SWIGLUOAI || act == ACT_SWIGLU_CLAMP; - const float up_bias = act == ACT_SWIGLUOAI ? 1.0f : 0.0f; - const float lim = swiglu_limit, alpha = swiglu_alpha; - for (int i = i0; i < i1; ++i) { - // gate = row i, up = row I+i - 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; - if (clamped) { - // clamp(gate, max=lim) * sigmoid(alpha * gate) * (clamp(up, +-lim) + up_bias) - // -- swigluoai carries the +1 up bias (gpt-oss/MiniMax); swiglu_clamp - // (GLM-5.3) does not. lim == +inf: no clamp. - 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 + up_bias)); - } else { - g_row[i] = f32_to_bf16(act_apply(act, gate) * up); - } - } - } - - void do_pass2(const MoeTask* t, int64_t p) { - if (fmt == WF_MXFP4) { - do_pass2_mxfp4(t, p); - return; - } - if (fmt == WF_DSFP4) { - do_pass2_dsfp4(t, p); - return; - } - const int64_t hb = p % n_hblk; - const int tok = static_cast(p / n_hblk); - const int h0 = static_cast(hb) * HBLK; - const int h1 = std::min(H, h0 + HBLK); - // Resolve this task's layer base once; row indexing below is layer-local (e). - 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)); - bf16_t* y_row = t->y + (size_t)tok * H; - for (int h = h0; h < h1; ++h) { - float acc = 0.0f; - for (int k = 0; k < top_k; ++k) { - const int e = t->ids[static_cast(tok) * top_k + k]; - if (e < 0 || e >= num_experts) continue; - const float w_out = apply_on_input ? 1.0f : t->w[static_cast(tok) * top_k + k]; - const size_t gr = (size_t)tok * top_k + k; - 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 += gemm2_dot(down_l, dn_packed_l, dn_scale_l, dn_global_l, e, h, g_row, ge, go, gi8, - gas) * w_out; - } - y_row[h] = f32_to_bf16(acc); - } - } - - // ----------------------------- mxfp4 (gpt-oss) ----------------------------- - // Transposed split-K layout (N innermost), so the GEMV streams K and accumulates - // a contiguous block of N output columns -> cache-line-efficient, no repack, no - // extra host memory. Pass 1 fuses gate_up + clamped-swiglu(+bias); pass 2 fuses - // down(+bias) * router-weight, summed over the token's routes. - // - // 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. - - void do_pass1_mxfp4(const MoeTask* t, int64_t p) { - const int64_t ib = p % n_iblk; - const int64_t tk = p / n_iblk; - const int k = static_cast(tk % top_k); - const int tok = static_cast(tk / top_k); - const int e = t->ids[static_cast(tok) * top_k + k]; - if (e < 0 || e >= num_experts) return; - // Resolve this task's layer base once; row indexing below is layer-local (e). - 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; // gate_up output width (gate/up interleaved) - const int Hh = H / 2; // packed-K rows - const int i0 = static_cast(ib) * IBLK; - const int i1 = std::min(I, i0 + IBLK); - const int nunit = i1 - i0; // intermediate units owned by this tile - const int col0 = 2 * i0; // first gate_up column - const int ncol = 2 * nunit; // gate_up columns owned by this tile - const bf16_t* x_row = t->x + (size_t)tok * H; - 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; - float gu[2 * IBLK]; - mxgemv(gu, blk_e + col0, scl_e + col0, x_row, Hh, N2, ncol, e2m1_lut, e8m0_lut); - const bf16_t* bias_e = gu_bias_l + (size_t)e * N2 + col0; - bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I; - const float lim = swiglu_limit, alpha = swiglu_alpha; - 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)); // gate * sigmoid(alpha*gate) - g_row[i0 + j] = f32_to_bf16(glu * (up + 1.0f)); - } + const float w_in = apply_on_input ? io_w[static_cast(tok) * top_k + k] : 1.0f; + const int i0 = static_cast(ib) * kRowBlock; + const int i1 = std::min(I, i0 + kRowBlock); + fmt->gate_up(static_cast(t.layer), e, tok, i0, i1, io_x + (size_t)tok * H, w_in, + g_scratch.data() + ((size_t)tok * top_k + k) * I); } - void do_pass2_mxfp4(const MoeTask* t, int64_t p) { + // Pass 2 item p = one kRowBlock block of one token's output row, summed over its routes. + void do_pass2(const MoeTask& t, int64_t p) { + const int top_k = params.top_k, H = params.H, I = params.I; const int64_t hb = p % n_hblk; const int tok = static_cast(p / n_hblk); - const int Ih = I / 2; // packed-K rows - const int h0 = static_cast(hb) * HBLK; - const int h1 = std::min(H, h0 + HBLK); - const int nh = h1 - h0; - // Resolve this task's layer base once; row indexing below is layer-local (e). - 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 = reinterpret_cast(tbl_at(dn_bias_tbl, t->layer_id)); - float acc[HBLK]; - for (int c = 0; c < nh; ++c) acc[c] = 0.0f; + const int h0 = static_cast(hb) * kRowBlock; + const int h1 = std::min(H, h0 + kRowBlock); + DownRoute routes[kMaxTopK]; + float w_out[kMaxTopK]; + int n = 0; for (int k = 0; k < top_k; ++k) { - const int e = t->ids[static_cast(tok) * top_k + k]; + const int e = io_ids[static_cast(tok) * top_k + k]; if (e < 0 || e >= num_experts) continue; - const float wt = t->w[static_cast(tok) * top_k + k]; - const bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I; - 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)(I / 32) * H; - float part[HBLK]; - mxgemv(part, blk_e + h0, scl_e + h0, g_row, 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[c] += (part[c] + bf16_to_f32(bias_e[c])) * wt; - } - 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[c]); - } - - // ----------------------------- ds_fp4 (DSV4) ------------------------------- - // Row-major e2m1 + e8m0/32 (no global); silu-swiglu with clamp; FP8-roundtripped - // activations (x once in submit -> xq_scratch; the intermediate g in a dedicated - // round-trip phase between the two passes). Router weight applies on the down output. - - void do_pass1_dsfp4(const MoeTask* t, int64_t p) { - const int64_t ib = p % n_iblk; - const int64_t tk = p / n_iblk; - const int k = static_cast(tk % top_k); - const int tok = static_cast(tk / top_k); - const int e = t->ids[static_cast(tok) * top_k + k]; - if (e < 0 || e >= num_experts) return; - // Resolve this task's layer base once; row indexing below is layer-local (e). - 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; - // fp8-roundtripped input, pre-deinterleaved to fp32 even/odd halves. - const float* xe = xe_scratch.data() + (size_t)tok * (H / 2); - const float* xo = xo_scratch.data() + (size_t)tok * (H / 2); - const uint8_t* gp = gu_packed_l + (size_t)e * N2 * Hh; - const uint8_t* gs = gu_scale_l + (size_t)e * N2 * Hs; - bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I; - 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) { - // gate_up is stored bf16 by the reference GEMV before swiglu; round to match. - 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)); // silu(gate) - g_row[i] = f32_to_bf16(glu * up); - } - } - - // Prepare one intermediate row (token,route) for the down GEMV, reused across every down output row. - // ds_fp4 first FP8 round-trips it (DSV4 act_quant); nvfp4/ds_fp4/fp8_block deinterleave it to fp32 even/odd; q4_0 Q8_0-quantizes it. - void prep_g_row(int64_t r) { - bf16_t* g = g_scratch.data() + (size_t)r * I; - if (use_q4a8) { // q4_0 W4A8: Q8_0-quantize the intermediate row for the down GEMV. - quant_q8_0(g, I, gi8_scratch.data() + (size_t)r * I, - gas_scratch.data() + (size_t)r * (I / 32)); - return; - } - if (fmt == WF_DSFP4) fp8_roundtrip_bf16(g, g, I, act_block); - float* ge = ge_scratch.data() + (size_t)r * (I / 2); - float* go = go_scratch.data() + (size_t)r * (I / 2); - deinterleave_bf16_f32(g, ge, go, I); - if (use_vnni) - quant_i8_pg16(ge, go, I, gi8_scratch.data() + (size_t)r * I, - gas_scratch.data() + (size_t)r * (I / 16)); - } - - void do_pass2_dsfp4(const MoeTask* t, int64_t p) { - const int64_t hb = p % n_hblk; - const int tok = static_cast(p / n_hblk); - const int Ih = I / 2, Is = I / 32; - const int h0 = static_cast(hb) * HBLK; - const int h1 = std::min(H, h0 + HBLK); - // Resolve this task's layer base once; row indexing below is layer-local (e). - 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)); - bf16_t* y_row = t->y + (size_t)tok * H; + const size_t r = (size_t)tok * top_k + k; + routes[n] = {e, static_cast(r), g_scratch.data() + r * I}; + w_out[n] = apply_on_input ? 1.0f : io_w[r]; + ++n; + } + float part[kMaxTopK * kRowBlock]; + if (n > 0) fmt->down(static_cast(t.layer), h0, h1, routes, n, part); + bf16_t* y_row = io_y + (size_t)tok * H; for (int h = h0; h < h1; ++h) { float acc = 0.0f; - for (int k = 0; k < top_k; ++k) { - const int e = t->ids[static_cast(tok) * top_k + k]; - if (e < 0 || e >= num_experts) continue; - const float wt = t->w[static_cast(tok) * top_k + k]; - const float* ge = ge_scratch.data() + ((size_t)tok * top_k + k) * (I / 2); - const float* go = go_scratch.data() + ((size_t)tok * top_k + k) * (I / 2); - const uint8_t* dp = dn_packed_l + (size_t)e * (size_t)H * Ih + (size_t)h * Ih; - const uint8_t* ds = dn_scale_l + (size_t)e * (size_t)H * Is + (size_t)h * Is; - // The reference rounds each route's weighted down output to bf16 before the - // fp32 sum over routes (down [T, top_k, H] bf16 -> .sum(dim=1)). - acc += bf16_to_f32(f32_to_bf16(dsdot(dp, ds, ge, go, I, e2m1_lut, e8m0_lut) * wt)); + for (int j = 0; j < n; ++j) { + const float v = part[j * kRowBlock + h - h0] * w_out[j]; + acc += rounds_routes ? bf16_to_f32(f32_to_bf16(v)) : v; } y_row[h] = f32_to_bf16(acc); } } - void run_task_body(const MoeTask* t) { + void run_task_body(const MoeTask& t) { int local_sense = 0; for (;;) { int64_t p = p1_next.fetch_add(1, std::memory_order_relaxed); @@ -2009,13 +282,13 @@ struct CpuMoeExecutor { do_pass1(t, p); } barrier(local_sense); - // Prepare the intermediate rows (per token,route) for the down GEMV, see prep_g_row. + // Prepare the intermediate rows (per token,route) for the down GEMV, see prep_inter. // Needs all of pass1 done (a full row spans every iblk). - if (needs_di || use_q4a8) { + if (preps_inter) { for (;;) { int64_t r = prt_next.fetch_add(1, std::memory_order_relaxed); if (r >= prt_total) break; - prep_g_row(r); + fmt->prep_inter(r, g_scratch.data() + (size_t)r * params.I); } barrier(local_sense); } @@ -2051,8 +324,7 @@ struct CpuMoeExecutor { } // submit() writes cur_task before the release store to `submitted`. my_gen = gen; - MoeTask* t = cur_task; - run_task_body(t); + run_task_body(cur_task); last_active = wclk::now(); if (done_count.fetch_add(1) + 1 == num_threads) { completed.store(my_gen, std::memory_order_release); @@ -2064,69 +336,20 @@ struct CpuMoeExecutor { } } - void submit(MoeTask* 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 - // construction-time hint (CUDA-graph capture warms the largest bs first, so - // 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; - prt_total = (needs_di || use_q4a8) ? static_cast(t->num_tokens) * top_k : 0; + void submit(MoeTask t) { + const int num_tokens = static_cast(t.num_tokens); + p1_total = static_cast(num_tokens) * params.top_k * n_iblk; + p2_total = static_cast(num_tokens) * n_hblk; + prt_total = preps_inter ? static_cast(num_tokens) * params.top_k : 0; p1_next.store(0, std::memory_order_relaxed); p2_next.store(0, std::memory_order_relaxed); prt_next.store(0, std::memory_order_relaxed); done_count.store(0, std::memory_order_relaxed); bar_count.store(0, std::memory_order_relaxed); bar_sense.store(0, std::memory_order_relaxed); - // Deinterleave the per-token input once, up front (ds_fp4 FP8-round-trips it first). + // Prepare the per-token input once, up front. // Single-threaded: tiny for decode, and done before the workers are woken below. - if (needs_di) { - const size_t xn = static_cast(t->num_tokens) * H; - if (xn / 2 > xe_scratch.size()) { - xe_scratch.resize(xn / 2); - xo_scratch.resize(xn / 2); - ge_scratch.resize(static_cast(t->num_tokens) * top_k * (I / 2)); - go_scratch.resize(static_cast(t->num_tokens) * top_k * (I / 2)); - if (use_vnni) { - xi8_scratch.resize(xn); - xas_scratch.resize(xn / 16); - gi8_scratch.resize(static_cast(t->num_tokens) * top_k * I); - gas_scratch.resize(static_cast(t->num_tokens) * top_k * (I / 16)); - } - } - const bool ds = (fmt == WF_DSFP4) && !input_prequant; - if (ds && xn > xq_scratch.size()) xq_scratch.resize(xn); - for (int tok = 0; tok < t->num_tokens; ++tok) { - const bf16_t* src = t->x + (size_t)tok * H; - if (ds) { // DSV4 FP8-round-trips the input before the gate_up GEMV - bf16_t* xq = xq_scratch.data() + (size_t)tok * H; - fp8_roundtrip_bf16(src, xq, H, act_block); - src = xq; - } - float* xe = xe_scratch.data() + (size_t)tok * (H / 2); - float* xo = xo_scratch.data() + (size_t)tok * (H / 2); - deinterleave_bf16_f32(src, xe, xo, H); - if (use_vnni) - quant_i8_pg16(xe, xo, H, xi8_scratch.data() + (size_t)tok * H, - xas_scratch.data() + (size_t)tok * (H / 16)); - } - } - // q4_0 W4A8: Q8_0-quantize the per-token input once (single-threaded, tiny for decode). - if (use_q4a8) { - const size_t xn = static_cast(t->num_tokens) * H; - if (xn > xi8_scratch.size()) { - xi8_scratch.resize(xn); - xas_scratch.resize(xn / 32); - gi8_scratch.resize(static_cast(t->num_tokens) * top_k * I); - gas_scratch.resize(static_cast(t->num_tokens) * top_k * (I / 32)); - } - for (int tok = 0; tok < t->num_tokens; ++tok) - quant_q8_0(t->x + (size_t)tok * H, H, xi8_scratch.data() + (size_t)tok * H, - xas_scratch.data() + (size_t)tok * (H / 32)); - } + for (int tok = 0; tok < num_tokens; ++tok) fmt->prep_input(tok, io_x + (size_t)tok * params.H); { std::lock_guard lk(task_mtx); cur_task = t; @@ -2153,62 +376,45 @@ struct CpuMoeExecutor { sync_cv.wait(lk, [&] { return completed.load(std::memory_order_acquire) >= target; }); } - void submit_with_cuda_stream(uintptr_t stream, uintptr_t task) { - cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::submit_cb, - reinterpret_cast(task)); - } - - void sync_with_cuda_stream(uintptr_t stream, uintptr_t task) { - cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::sync_cb, - reinterpret_cast(task)); - } - - // Register a (layer, batch-size) slot's task so the coordinator can dispatch it on a - // flag bump. - void register_flag_task(int slot, uintptr_t task) { - std::lock_guard lk(flag_task_mtx); - if (static_cast(flag_task.size()) <= slot) flag_task.resize(slot + 1, nullptr); - flag_task[slot] = reinterpret_cast(task); + void submit_with_cuda_stream(uintptr_t stream, int layer_id, int num_tokens) { + const MoeTask t = task(layer_id, num_tokens); + HostTask& h = host_tasks.try_emplace(to_word(t), HostTask{this, t}).first->second; + cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::submit_cb, &h); } - int64_t flag_served_count(int slot) const { - return (slot >= 0 && slot < static_cast(flag_served.size())) ? flag_served[slot] : 0; + void sync_with_cuda_stream(uintptr_t stream) { + cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::sync_cb, this); } - // Start the busy-poll coordinator over the mapped-pinned flag arrays. ``pin_core`` >= 0 + // Start the busy-poll coordinator over the doorbell. ``pin_core`` >= 0 // pins the coordinator to that logical CPU (the worker auto-sizing reserves it), so its // polling never migrates onto / contends with a GEMV worker's core. - void start_flag_coordinator(uintptr_t ready_ptr, uintptr_t done_ptr, int num_slots, - int pin_core) { - ready_flags = reinterpret_cast(ready_ptr); - done_flags = reinterpret_cast(done_ptr); - coord_num_slots = num_slots; - { - std::lock_guard lk(flag_task_mtx); - if (static_cast(flag_task.size()) < num_slots) flag_task.resize(num_slots, nullptr); - } - flag_served.assign(num_slots, 0); + void start_flag_coordinator(int pin_core) { + gpu_sync = std::make_unique(); coord_stop.store(false); coord_pinned = CPU_MOE_HAS_AFFINITY && pin_core >= 0; coord_thread = std::thread([this, pin_core] { -#if CPU_MOE_HAS_AFFINITY - if (pin_core >= 0) { - cpu_set_t set; - CPU_ZERO(&set); - CPU_SET(pin_core, &set); - pthread_setaffinity_np(pthread_self(), sizeof(set), &set); - } -#endif + if (pin_core >= 0) pin_this_thread(pin_core); coordinator_loop(); }); } + void flag_signal(uintptr_t stream, int layer_id, int num_tokens) { + gpu_sync->gpu_submit(stream, to_word(task(layer_id, num_tokens))); + } + void flag_wait(uintptr_t stream) { gpu_sync->gpu_wait(stream); } + bool flag_pending() const { return gpu_sync && gpu_sync->pending(); } + // Complete a doorbell the coordinator never answered so the stream can make progress. + void poison_flag() { + if (gpu_sync) gpu_sync->cpu_done(); + } + void coordinator_loop() { // Idle backoff (edge-device friendliness): spin hot only while decode traffic is - // flowing, with a TIME-based hot window (an iteration count varies 2-6x with slot - // count and pause cost, so some configs dozed every token). 50 ms since the last - // flag comfortably covers intra- and inter-token gaps -- including host-func-only - // stretches (prefill bursts, --moe-cpu-layers subsets) -- so steady decode never + // flowing, with a TIME-based hot window (an iteration count varies 2-6x with pause + // cost, so some configs dozed every token). 50 ms since the last + // flag comfortably covers intra- and inter-token gaps -- including prefill bursts and + // --moe-cpu-layers gaps -- so steady decode never // sleeps and keeps the ~us-level wakeup. Past it, i.e. the engine is actually idle, // escalate to timed sleeps capped at 2 ms: a dozing coordinator costs <0.1% of a // core instead of 100%, and the only price is a <=2 ms discovery delay on the FIRST @@ -2222,30 +428,13 @@ struct CpuMoeExecutor { int64_t sleep_us = 100; bool dozing = false; while (!coord_stop.load(std::memory_order_relaxed)) { - bool any = false; - for (int L = 0; L < coord_num_slots; ++L) { - // Binary handshake (memop-compatible: the GPU-side WAIT compares against an - // immediate baked at graph capture, so the protocol resets per step instead of - // counting). Acquire: everything the GPU made visible before setting ready -- - // the D2H input copies -- is visible to the worker pool after this read. - if (flag_load_acquire(&ready_flags[L]) != 0) { - flag_store_release(&ready_flags[L], 0); // consume this step's doorbell - MoeTask* t; - { - std::lock_guard lk(flag_task_mtx); - t = (L < static_cast(flag_task.size())) ? flag_task[L] : nullptr; - } - if (t != nullptr) { - submit(t); - sync(coord_pinned && spin_wait.load(std::memory_order_relaxed)); - } - // Release: the workers' y stores are visible before the GPU sees done. - flag_store_release(&done_flags[L], 1); - if (L < static_cast(flag_served.size())) ++flag_served[L]; - any = true; - } - } - if (any) { + const uint64_t word = gpu_sync->cpu_poll(); + if (word != 0) { + submit(from_word(word)); + sync(coord_pinned && spin_wait.load(std::memory_order_relaxed)); + // Release: the workers' y stores are visible before the GPU sees done. + gpu_sync->cpu_done(); + flag_served.fetch_add(1, std::memory_order_relaxed); last_active = coord_clock::now(); empty_polls = 0; sleep_us = 100; @@ -2254,15 +443,13 @@ struct CpuMoeExecutor { } // Hot phase: pause-spin, consulting the clock only every 1024 empty polls (the // amortization must gate the CLOCK, not the sleep -- gating the sleep decision - // let 1023/1024 idle iterations hot-scan all slots, ~20% of a core at 992 slots). + // let 1023/1024 idle iterations hot-poll). // Doze phase: sleep on EVERY iteration until traffic returns; one wake-up scan // (~0.5 us) per 2 ms sleep is ~0.03% duty. if (!dozing) { if ((++empty_polls & 1023u) != 0 || coord_clock::now() - last_active < kHotWindow) { -#if CPU_MOE_X86 - _mm_pause(); -#endif + cpu_relax(); continue; } dozing = true; @@ -2270,92 +457,73 @@ struct CpuMoeExecutor { std::this_thread::sleep_for(std::chrono::microseconds(sleep_us)); sleep_us = std::min(sleep_us * 2, kSleepCapUs); } - // Teardown: release any in-flight (or future) spin-wait immediately so a replay - // caught mid-shutdown exits its sync kernel now instead of owning the watchdog - // stall. Runs before the destructor's join() returns, while the flag arrays are - // still alive on the Python side. - for (int L = 0; L < coord_num_slots; ++L) { - flag_store_release(&done_flags[L], INT64_MAX); - } + // Teardown: complete any in-flight wait so a replay caught mid-shutdown does not stall + // until the watchdog fires. Runs before the destructor's join() returns, while the doorbell is still alive. + gpu_sync->cpu_done(); } // Eager (non-graph) path: run one task to completion on the pool. - void run_task(uintptr_t task) { - MoeTask* t = reinterpret_cast(task); - submit(t); + void run(int layer_id, int num_tokens) { + submit(task(layer_id, num_tokens)); sync(); } static void CUDART_CB submit_cb(void* ud) { - MoeTask* t = reinterpret_cast(ud); - t->exec->submit(t); - } - static void CUDART_CB sync_cb(void* ud) { - MoeTask* t = reinterpret_cast(ud); - t->exec->sync(); + HostTask* h = reinterpret_cast(ud); + h->exec->submit(h->task); } + static void CUDART_CB sync_cb(void* ud) { reinterpret_cast(ud)->sync(); } }; } // namespace +} // namespace freetoken::cpu_moe PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { namespace py = pybind11; + using freetoken::cpu_moe::CpuMoeExecutor; + using freetoken::cpu_moe::GpuSync; py::class_(m, "CpuMoeExecutor") - .def(py::init, int, int>(), + .def(py::init, std::map, + double, double, std::vector, uintptr_t, uintptr_t, uintptr_t, uintptr_t>(), py::arg("num_threads"), py::arg("num_layers"), py::arg("num_experts"), py::arg("top_k"), py::arg("hidden_size"), py::arg("inter_size"), py::arg("max_tokens"), py::arg("activation_id"), py::arg("apply_router_weight_on_input"), py::arg("weight_format"), - py::arg("gate_up_ptr"), py::arg("down_ptr"), py::arg("gate_up_scale_ptr"), - py::arg("gate_up_global_ptr"), py::arg("down_scale_ptr"), - py::arg("down_global_ptr"), py::arg("gate_up_bias_ptr"), - py::arg("down_bias_ptr"), py::arg("swiglu_alpha"), py::arg("swiglu_limit"), - py::arg("core_ids"), py::arg("fp8_block_gu_scale_stride") = 0, - py::arg("fp8_block_dn_scale_stride") = 0) - .def("create_task", &CpuMoeExecutor::create_task, py::arg("layer_id"), - py::arg("num_tokens"), py::arg("x_ptr"), py::arg("ids_ptr"), py::arg("w_ptr"), + py::arg("banks"), py::arg("options"), py::arg("swiglu_alpha"), py::arg("swiglu_limit"), + py::arg("core_ids"), py::arg("x_ptr"), py::arg("ids_ptr"), py::arg("w_ptr"), py::arg("y_ptr")) .def("submit_with_cuda_stream", &CpuMoeExecutor::submit_with_cuda_stream, - py::arg("stream"), py::arg("task"), py::call_guard()) + py::arg("stream"), py::arg("layer_id"), py::arg("num_tokens"), + py::call_guard()) .def("sync_with_cuda_stream", &CpuMoeExecutor::sync_with_cuda_stream, - py::arg("stream"), py::arg("task"), py::call_guard()) - .def("run_task", &CpuMoeExecutor::run_task, py::arg("task"), + py::arg("stream"), py::call_guard()) + .def("run", &CpuMoeExecutor::run, py::arg("layer_id"), py::arg("num_tokens"), py::call_guard()) - .def("register_flag_task", &CpuMoeExecutor::register_flag_task, - py::arg("slot"), py::arg("task")) - .def("flag_served_count", &CpuMoeExecutor::flag_served_count, py::arg("slot")) - .def("start_flag_coordinator", &CpuMoeExecutor::start_flag_coordinator, - py::arg("ready_ptr"), py::arg("done_ptr"), py::arg("num_slots"), - py::arg("pin_core")) - .def("set_input_prequant", - [](CpuMoeExecutor& e, bool v) { e.input_prequant = v; }, - py::arg("value")) - .def("set_act_block", - [](CpuMoeExecutor& e, int v) { e.act_block = v; }, - py::arg("value")) + .def("flag_served", [](const CpuMoeExecutor& e) { return e.flag_served.load(); }) + .def("start_flag_coordinator", &CpuMoeExecutor::start_flag_coordinator, py::arg("pin_core")) + .def("flag_signal", &CpuMoeExecutor::flag_signal, py::arg("stream"), py::arg("layer_id"), + py::arg("num_tokens")) + .def("flag_wait", &CpuMoeExecutor::flag_wait, py::arg("stream")) + .def("flag_pending", &CpuMoeExecutor::flag_pending) + .def("poison_flag", &CpuMoeExecutor::poison_flag) .def("isa_name", &CpuMoeExecutor::isa_name) - .def("fp8_block_isa_name", &CpuMoeExecutor::fp8_block_isa_name) .def("set_spin_wait", [](CpuMoeExecutor& e, bool v) { e.spin_wait.store(v, std::memory_order_relaxed); }, py::arg("enabled")) .def("get_spin_wait", [](const CpuMoeExecutor& e) { return e.spin_wait.load(); }); - m.def("memops_probe", &cumemops_probe, py::arg("stream"), py::arg("scratch_addr")); - m.def("memop_submit", &cumemop_submit, py::arg("stream"), py::arg("done_addr"), - py::arg("ready_addr"), py::arg("slot")); - m.def("memop_sync", &cumemop_sync, py::arg("stream"), py::arg("done_addr"), - py::arg("slot")); + m.def("memops_probe", &GpuSync::probe, py::arg("stream")); + m.def("host_func_capture_safe", &GpuSync::host_func_capture_safe); // ABI capability marker: the highest ActKind this build implements in the // GENERIC epilogue. CpuMoeExecutor.__init__ probes it before requesting an act // id the epilogue must handle -- a prebuilt .so from before ACT_SWIGLUOAI // accepts id 3 without error and silently computes the wrong activation // (act_apply falls through to gelu_tanh); the probe turns a stale extension // into a loud rebuild instruction instead of wrong model outputs. - m.def("max_generic_act_id", []() { return static_cast(ACT_SWIGLU_CLAMP); }); - // Weight-format ABI probe. A stale prebuilt .so accepts a newer WFmt id without - // complaint and then falls through to the wrong dequant branch -- silently wrong - // numbers, not a crash. The Python executor refuses to build against an extension - // that predates the format it was asked for. - m.def("max_weight_format_id", []() { return static_cast(WF_FP8_BLOCK); }); + m.def("max_generic_act_id", []() { return static_cast(freetoken::cpu_moe::ACT_SWIGLU_CLAMP); }); + m.def("weight_formats", []() { + py::list names; + for (const auto& f : freetoken::cpu_moe::kFormats) names.append(f.first); + return names; + }); } diff --git a/python/freetoken/kernel/csrc/cpu_moe/format.h b/python/freetoken/kernel/csrc/cpu_moe/format.h new file mode 100644 index 0000000000..286ae110f7 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/format.h @@ -0,0 +1,131 @@ +// The interface between the CPU MoE executor and one expert weight format. +#pragma once + +#include +#include +#include +#include + +#include "common.h" + +namespace freetoken::cpu_moe { + +// Output-row tiling. Small enough to give every worker independent work even at +// batch size 1; large enough to amortize the atomic work-grab. +// +// Bandwidth notes (Sapphire Rapids 8480+, 13 cores): the two passes already read +// every expert weight byte exactly once per token (each output row block is owned +// by one worker), and x stays hot in L1 across a (token,expert)'s rows -- so the +// kernel is single-read bandwidth-optimal at bs=1 (~205 GB/s vs ~55 GB/s PCIe). +// One worker per *physical* core, pinned, is the sweet spot; SMT oversubscription +// thrashes the spin-barrier. Deferred (not worth it here / for this workload): +// - AMX-bf16: a GEMM tile engine; decode is M=1 GEMV so tiles sit idle. It would +// only pay off in a grouped/batched (dedup) path. +// - expert dedup for bs>1: read each distinct expert once and GEMM its tokens. +// Helps locality+bytes when bs is large; decode batches here are tiny (<=4). +// - NUMA: a single node is assumed. Multi-socket machines would split each +// expert's K dimension per node (banks are already per-row contiguous). +constexpr int kRowBlock = 32; + +struct MoeParams { + int top_k = 0, max_tokens = 0; + int H = 0, I = 0; + int act = ACT_SILU; + float swiglu_alpha = 1.0f; + float swiglu_limit = INFINITY; // +inf == no clamp +}; + +// What cpu_executor.py passes for one format, by name: banks map canonical roles (gate_up, gate_up_scale, ...) to per-layer tables, options hold small integers. +// A format reads what it needs once, in its constructor. +struct FormatArgs { + std::string format; + std::map banks; + std::map options; + + const uint64_t* bank(const std::string& role) const { + const auto it = banks.find(role); + if (it == banks.end() || it->second == 0) + throw std::invalid_argument("CPU MoE " + format + " experts need the " + role + " bank"); + return reinterpret_cast(it->second); + } + int64_t option(const std::string& name) const { + const auto it = options.find(name); + if (it == options.end()) throw std::invalid_argument("CPU MoE " + format + " experts need the " + name + " option"); + return it->second; + } + int64_t option(const std::string& name, int64_t fallback) const { + const auto it = options.find(name); + return it == options.end() ? fallback : it->second; + } +}; + +// A bank table is the address of a CPU int64 array of length +// num_layers (one base address per layer, built by cpu_executor.py's +// _make_table), not a single flat bank. tbl_at resolves +// tbl[layer_id] once per task/pass. +inline const void* tbl_at(const uint64_t* tbl, int layer_id) { + return reinterpret_cast(tbl[layer_id]); +} + +// One route of a pass 2 item: expert e on intermediate row r = tok * top_k + k, whose bf16 values are g. +struct DownRoute { + int e; + int64_t r; + const bf16_t* g; +}; + +// g[i] for i in [i0, i1) from the gate row i and the up row I + i; dot(row) is the row's GEMV. +// ACT_SWIGLUOAI and ACT_SWIGLU_CLAMP clamp both inputs; the other kinds go through act_apply. +template +inline void gated_rows(const MoeParams& p, int i0, int i1, float w_in, bf16_t* g, RowDot dot) { + const bool clamped = p.act == ACT_SWIGLUOAI || p.act == ACT_SWIGLU_CLAMP; + const float up_bias = p.act == ACT_SWIGLUOAI ? 1.0f : 0.0f; + const float lim = p.swiglu_limit, alpha = p.swiglu_alpha; + for (int i = i0; i < i1; ++i) { + // gate = row i, up = row I+i + float gate = dot(i) * w_in; + float up = dot(p.I + i) * w_in; + if (clamped) { + // clamp(gate, max=lim) * sigmoid(alpha * gate) * (clamp(up, +-lim) + up_bias) + // -- swigluoai carries the +1 up bias (gpt-oss/MiniMax); swiglu_clamp + // (GLM-5.3) does not. lim == +inf: no clamp. + 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[i] = f32_to_bf16(glu * (up + up_bias)); + } else { + g[i] = f32_to_bf16(act_apply(p.act, gate) * up); + } + } +} + +// One expert weight format: its kernels, its activation layout and the scratch for it. +// The executor owns the threads and the work split. +// Scratch is sized once from max_tokens; gate_up, down and prep_inter run on the worker threads, prep_input before they wake. +class MoeFormat { + public: + explicit MoeFormat(const MoeParams& p) : p(p) {} + virtual ~MoeFormat() = default; + + virtual std::string isa_name() const = 0; + // Once per input row x[tok] before pass 1. + virtual void prep_input(int /*tok*/, const bf16_t* /*x*/) {} + // Whether pass 2 needs prep_inter on every intermediate row. + virtual bool preps_inter() const { return false; } + virtual void prep_inter(int64_t /*r*/, bf16_t* /*g*/) {} + // Pass 1: g[i] for i in [i0, i1) (at most kRowBlock rows), expert e of `layer` on token tok. + // w_in scales the gate and up dots: the router weight when it applies on the input, else 1. + virtual void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* x, float w_in, + bf16_t* g) const = 0; + // Pass 2: out[j * kRowBlock + h - h0] for h in [h0, h1) (at most kRowBlock rows) and each of the n routes, before the router weight. + // Visiting the routes' rows h together keeps n weight streams in flight, which the down GEMV needs to reach DRAM bandwidth. + virtual void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const = 0; + // Round each router-weighted down output to bf16 before the sum over routes. + virtual bool rounds_routes() const { return false; } + + protected: + const MoeParams& p; +}; + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h new file mode 100644 index 0000000000..bf58874168 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h @@ -0,0 +1,72 @@ +// GPU <-> CPU handshake for the CPU MoE coordinator, one implementation per GPU runtime (gpu_sync_cuda.cpp, gpu_sync_hip.cpp). +// +// A backend keeps two contracts, in eager mode and in graph replay: +// 1. submit, GPU -> CPU: after gpu_submit(stream, task), cpu_poll() returns task once, and every earlier write of the stream is visible to the CPU. +// 2. wait, CPU -> GPU: gpu_wait(stream) holds later work of the stream until cpu_done(), and that work sees every CPU write made before cpu_done(). +// task is a nonzero word, opaque here; the executor encodes the layer and token count in it, as a constant baked into each captured graph. +// One task is in flight at a time: decode runs on one stream and every layer waits before the next one signals. +#pragma once + +#include +#include + +#if defined(_MSC_VER) +#include +#endif + +namespace freetoken::cpu_moe { + +// Portable ordering for the flag handshake: "ready observed => the DMA'd inputs that +// preceded the bump are visible" and "y stores are visible before done". Plain +// volatile loads lean on x86 TSO; acquire/release makes it hold on aarch64 too +// (GH200/Jetson) at zero x86 cost. (MSVC branch is x86-only today: compiler barrier +// + TSO.) +inline int64_t flag_load_acquire(const volatile int64_t* p) { +#if defined(_MSC_VER) + const int64_t v = *p; + _ReadWriteBarrier(); + return v; +#else + return __atomic_load_n(const_cast(p), __ATOMIC_ACQUIRE); +#endif +} + +inline void flag_store_release(volatile int64_t* p, int64_t v) { +#if defined(_MSC_VER) + _ReadWriteBarrier(); + *p = v; +#else + __atomic_store_n(const_cast(p), v, __ATOMIC_RELEASE); +#endif +} + +class GpuSync { + public: + // Allocates the handshake memory; throws if this runtime cannot. + GpuSync(); + ~GpuSync(); + GpuSync(const GpuSync&) = delete; + GpuSync& operator=(const GpuSync&) = delete; + + // Whether both contracts hold on this device and stream, including graph replay where the runtime needs a real check. + static bool probe(uintptr_t stream); + // Whether the cudaLaunchHostFunc fallback is replayed correctly from a captured graph. + static bool host_func_capture_safe(); + + // 1. submit. gpu_submit() is enqueued on `stream`, which may be capturing; cpu_poll() returns 0 when nothing is submitted. + void gpu_submit(uintptr_t stream, uint64_t task); + uint64_t cpu_poll(); + + // 2. wait. cpu_done() also unblocks the stream on teardown and from the watchdog. + void gpu_wait(uintptr_t stream); + void cpu_done(); + + // Submitted but not polled yet (watchdog only). + bool pending() const; + + private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp new file mode 100644 index 0000000000..0b73a91d61 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp @@ -0,0 +1,156 @@ +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#include +#else +#include +#endif + +#include + +#include "gpu_sync.h" + +namespace freetoken::cpu_moe { +namespace { + +// ===================================================================================== +// CUDA stream memory operations (driver API, resolved via dlopen -- no link-time or +// toolchain dependence). The GPU side of the flag handshake: submit = WRITE_VALUE +// (done=0 then ready=task word), sync = WAIT_VALUE(done >= 1). The wait is +// executed by the GPU front-end (no SM-resident kernel), so GPU "utilization" stays +// truthful during CPU compute windows -- a resident spin kernel pinned it at 99%, +// which laptop CPU/GPU dynamic power schedulers answered by clamping the CPU's max +// frequency (GEMV workers -1.5x: the reported edge regression). Availability is +// probed functionally at startup (memops_probe); anything unsupported (some driver models, +// vGPU, old drivers) falls back to the cudaLaunchHostFunc path. +#if defined(_WIN32) +static void* cumemop_dlopen() { return (void*)::LoadLibraryA("nvcuda.dll"); } +static void* cumemop_dlsym(void* h, const char* n) { + return (void*)::GetProcAddress((HMODULE)h, n); +} +#else +static void* cumemop_dlopen() { + void* h = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL); + if (h == nullptr) h = dlopen("libcuda.so", RTLD_LAZY | RTLD_LOCAL); + return h; +} +static void* cumemop_dlsym(void* h, const char* n) { return dlsym(h, n); } +#endif + +using cuMemOp64_fn = int (*)(void* stream, unsigned long long addr, unsigned long long value, + unsigned int flags); +static cuMemOp64_fn g_cu_write64 = nullptr; +static cuMemOp64_fn g_cu_wait64 = nullptr; +static constexpr unsigned int kCuWaitValueGeq = 0x0; // CU_STREAM_WAIT_VALUE_GEQ +static constexpr unsigned int kCuWriteDefault = 0x0; // CU_STREAM_WRITE_VALUE_DEFAULT + +static bool cumemop_resolve() { + static bool resolved = [] { + void* h = cumemop_dlopen(); + if (h == nullptr) return false; + // 11.7+ made the v2 entry points the default; older drivers export only the v1 + // names with the same signature. + g_cu_write64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWriteValue64_v2")); + if (g_cu_write64 == nullptr) + g_cu_write64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWriteValue64")); + g_cu_wait64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWaitValue64_v2")); + if (g_cu_wait64 == nullptr) + g_cu_wait64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWaitValue64")); + return g_cu_write64 != nullptr && g_cu_wait64 != nullptr; + }(); + return resolved; +} + +// GPU side of the flag handshake (see the block comment above): enqueued on the +// caller's (possibly capturing) stream; the WAIT immediate is the constant 1, +// replay-safe under CUDA graphs. +// The startup probe validates EAGER memops; a driver could still reject them at graph +// capture time. Those enqueue errors would otherwise be swallowed here and surface only +// as a later EndCapture failure -- log the first CUresult so triage is one step. +static void cumemop_check(int rc, const char* what) { + static std::atomic warned{false}; + if (rc != 0 && !warned.exchange(true)) { + std::fprintf(stderr, + "[freetoken/cpu_moe] %s failed with CUresult=%d (first occurrence; " + "subsequent errors are not repeated). If this happened during CUDA " + "graph capture, the driver lacks capture support for stream memops -- " + "set FREETOKEN_CPU_MOE_FLAG_SYNC=0.\n", + what, rc); + } +} + +} // namespace + +// ready and done live on separate cache lines of one mapped-pinned block; with UVA the GPU memops and the coordinator use the same addresses. +struct GpuSync::Impl { + void* block = nullptr; + volatile int64_t* ready = nullptr; + volatile int64_t* done = nullptr; +}; + +GpuSync::GpuSync() : impl_(std::make_unique()) { + if (!cumemop_resolve()) throw std::runtime_error("CPU MoE flag sync: CUDA stream memops are unavailable"); + const cudaError_t rc = cudaHostAlloc(&impl_->block, 128, cudaHostAllocPortable | cudaHostAllocMapped); + if (rc != cudaSuccess) + throw std::runtime_error(std::string("CPU MoE flag allocation failed: ") + cudaGetErrorString(rc)); + std::memset(impl_->block, 0, 128); + impl_->ready = static_cast(impl_->block); + impl_->done = impl_->ready + 8; +} + +GpuSync::~GpuSync() { cudaFreeHost(impl_->block); } + +// Functional probe on a scratch pinned int64: enqueue WRITE(7) + WAIT(>=7) + sync. +// Returns true only if the whole memop path works on THIS stream/device/driver. +bool GpuSync::probe(uintptr_t stream) { + if (!cumemop_resolve()) return false; + void* scratch = nullptr; + if (cudaHostAlloc(&scratch, sizeof(int64_t), cudaHostAllocPortable | cudaHostAllocMapped) != cudaSuccess) + return false; + *static_cast(scratch) = 0; + auto* s = reinterpret_cast(stream); + const auto addr = reinterpret_cast(scratch); + const bool ok = g_cu_write64(s, addr, 7ULL, kCuWriteDefault) == 0 && g_cu_wait64(s, addr, 7ULL, kCuWaitValueGeq) == 0 && + cudaStreamSynchronize(reinterpret_cast(stream)) == cudaSuccess; + // a failed enqueue can leave the WRITE pending on the scratch word + cudaStreamSynchronize(reinterpret_cast(stream)); + cudaFreeHost(scratch); + return ok; +} + +bool GpuSync::host_func_capture_safe() { return true; } + +void GpuSync::gpu_submit(uintptr_t stream, uint64_t task) { + auto* s = reinterpret_cast(stream); + // Order matters and is preserved by the front end: reset done BEFORE raising ready, + // so the coordinator's completion write for THIS step can never be wiped. + cumemop_check(g_cu_write64(s, reinterpret_cast(impl_->done), 0ULL, kCuWriteDefault), + "cuStreamWriteValue64(done)"); + cumemop_check(g_cu_write64(s, reinterpret_cast(impl_->ready), task, kCuWriteDefault), + "cuStreamWriteValue64(ready)"); +} + +void GpuSync::gpu_wait(uintptr_t stream) { + cumemop_check(g_cu_wait64(reinterpret_cast(stream), reinterpret_cast(impl_->done), 1ULL, + kCuWaitValueGeq), + "cuStreamWaitValue64(done)"); +} + +uint64_t GpuSync::cpu_poll() { + // Acquire: everything the GPU made visible before writing ready -- the D2H input copies -- is visible to the worker pool after this read. + const int64_t task = flag_load_acquire(impl_->ready); + if (task == 0) return 0; + flag_store_release(impl_->ready, 0); // consume this step's doorbell + return static_cast(task); +} + +void GpuSync::cpu_done() { flag_store_release(impl_->done, 1); } + +bool GpuSync::pending() const { return flag_load_acquire(impl_->ready) != 0 && flag_load_acquire(impl_->done) == 0; } + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp new file mode 100644 index 0000000000..4baf76e768 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp @@ -0,0 +1,260 @@ +// HIP stream memory operations for the CPU MoE flag handshake. +// +// HIP WAIT addresses must come from hipMallocSignalMemory, one 64-bit word per allocation. +// +// ROCm 10 records hipStreamWrite/WaitValue64 into the graph under stream capture, one node per call. +// ROCm 7.14 runs them while the stream is being captured but records nothing, so there we add explicit batch-memop nodes and splice them into the stream's dependency set. +// Which of the two applies is measured once at runtime (records_native_memops), not read from the version number. +// The probe then does a real capture + instantiate + replay through the chosen path. +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "gpu_sync.h" + +namespace freetoken::cpu_moe { + +namespace { + +void hip_check(hipError_t rc, const char* what) { + if (rc != hipSuccess) + throw std::runtime_error(std::string("CPU MoE HIP flag sync: ") + what + " failed: " + hipGetErrorString(rc)); +} + +hipError_t alloc_signal(uint64_t** out) { + void* storage = nullptr; + const hipError_t rc = hipExtMallocWithFlags(&storage, sizeof(uint64_t), hipMallocSignalMemory); + if (rc == hipSuccess) { + *out = static_cast(storage); + __atomic_store_n(*out, 0ULL, __ATOMIC_RELEASE); + } + return rc; +} + +void free_signal(uint64_t* signal) { + if (signal != nullptr) (void)hipFree(signal); +} + +volatile int64_t* as_flag(uint64_t* signal) { return reinterpret_cast(signal); } + +bool capturing(hipStream_t stream) { + hipStreamCaptureStatus status = hipStreamCaptureStatusNone; + hip_check(hipStreamIsCapturing(stream, &status), "capture status query"); + return status == hipStreamCaptureStatusActive; +} + +// Whether this runtime records a native hipStreamWriteValue64 issued under capture as a graph node. +bool records_native_memops() { + static const bool native = [] { + uint64_t* word = nullptr; + hipStream_t stream = nullptr; + hipGraph_t graph = nullptr; + hipError_t rc = alloc_signal(&word); + if (rc == hipSuccess) rc = hipStreamCreateWithFlags(&stream, hipStreamNonBlocking); + if (rc == hipSuccess) rc = hipStreamBeginCapture(stream, hipStreamCaptureModeRelaxed); + const bool capture_started = rc == hipSuccess; + if (rc == hipSuccess) rc = hipStreamWriteValue64(stream, word, 1, hipStreamWriteValueDefault); + if (capture_started) { + const hipError_t end_rc = hipStreamEndCapture(stream, &graph); + if (rc == hipSuccess) rc = end_rc; + } + size_t nodes = 0; + if (rc == hipSuccess) rc = hipGraphGetNodes(graph, nullptr, &nodes); + if (graph != nullptr) (void)hipGraphDestroy(graph); + if (stream != nullptr) { + (void)hipStreamSynchronize(stream); + (void)hipStreamDestroy(stream); + } + free_signal(word); + return rc == hipSuccess && nodes == 1; + }(); + return native; +} + +// Appends one batch-memop node to the capture of `stream`; the node keeps a pointer to `ops`. +hipError_t add_capture_memop_node(hipStream_t stream, hipStreamBatchMemOpParams* ops, unsigned int count) { + hipStreamCaptureStatus status = hipStreamCaptureStatusNone; + hipGraph_t graph = nullptr; + const hipGraphNode_t* dependencies = nullptr; + size_t dependency_count = 0; + hipError_t rc = hipStreamGetCaptureInfo_v2(stream, &status, nullptr, &graph, &dependencies, &dependency_count); + if (rc != hipSuccess) return rc; + if (status != hipStreamCaptureStatusActive || graph == nullptr) return hipErrorIllegalState; + + hipCtx_t context = nullptr; +#if defined(__GNUC__) +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wdeprecated-declarations" +#endif + rc = hipCtxGetCurrent(&context); +#if defined(__GNUC__) +#pragma GCC diagnostic pop +#endif + if (rc != hipSuccess) return rc; + + hipBatchMemOpNodeParams params{}; + params.ctx = context; + params.count = count; + params.paramArray = ops; + hipGraphNode_t node = nullptr; + rc = hipGraphAddBatchMemOpNode(&node, graph, dependencies, dependency_count, ¶ms); + if (rc != hipSuccess) return rc; + return hipStreamUpdateCaptureDependencies(stream, &node, 1, hipStreamSetCaptureDependencies); +} + +hipStreamBatchMemOpParams write_op(uint64_t* address, uint64_t value) { + hipStreamBatchMemOpParams op{}; + op.writeValue.operation = hipStreamMemOpWriteValue64; + op.writeValue.address = reinterpret_cast(address); + op.writeValue.value64 = value; + return op; +} + +hipStreamBatchMemOpParams wait_op(uint64_t* address) { + hipStreamBatchMemOpParams op{}; + op.waitValue.operation = hipStreamMemOpWaitValue64; + op.waitValue.address = reinterpret_cast(address); + op.waitValue.value64 = 1; + op.waitValue.flags = hipStreamWaitValueGte; + return op; +} + +bool needs_explicit_nodes(hipStream_t stream) { return capturing(stream) && !records_native_memops(); } + +// node_op is null for a native call, else the parameter slot of an explicit node, which the node keeps a pointer to. +// The ops of one batch node run in parallel, so each write gets its own node and consecutive writes stay in order. +hipError_t write_value(hipStream_t stream, uint64_t* address, uint64_t value, hipStreamBatchMemOpParams* node_op) { + if (node_op == nullptr) return hipStreamWriteValue64(stream, address, value, hipStreamWriteValueDefault); + *node_op = write_op(address, value); + return add_capture_memop_node(stream, node_op, 1); +} + +hipError_t wait_done(hipStream_t stream, uint64_t* done, hipStreamBatchMemOpParams* node_op) { + if (node_op == nullptr) return hipStreamWaitValue64(stream, done, 1, hipStreamWaitValueGte, UINT64_MAX); + *node_op = wait_op(done); + return add_capture_memop_node(stream, node_op, 1); +} + +} // namespace + +struct GpuSync::Impl { + uint64_t* ready = nullptr; // signal memory + uint64_t* done = nullptr; + // Explicit nodes keep these addresses, so one parameter pair per distinct task, never moved (std::map nodes are stable). + std::map> submit_ops; // done = 0, ready = task + hipStreamBatchMemOpParams wait_op{}; +}; + +GpuSync::GpuSync() : impl_(std::make_unique()) { + hipError_t rc = alloc_signal(&impl_->ready); + if (rc == hipSuccess) rc = alloc_signal(&impl_->done); + if (rc != hipSuccess) { + free_signal(impl_->ready); + free_signal(impl_->done); + hip_check(rc, "signal allocation"); + } + // Measure now: the measurement captures and allocates, which a graph capture in progress would reject. + (void)records_native_memops(); +} + +GpuSync::~GpuSync() { + free_signal(impl_->ready); + free_signal(impl_->done); +} + +bool GpuSync::probe(uintptr_t /*stream*/) { + int device = 0; + int can_wait = 0; + if (hipGetDevice(&device) != hipSuccess || + hipDeviceGetAttribute(&can_wait, hipDeviceAttributeCanUseStreamWaitValue, device) != hipSuccess || can_wait == 0) + return false; + const bool explicit_nodes = !records_native_memops(); + + uint64_t* ready = nullptr; + uint64_t* done = nullptr; + hipStream_t stream = nullptr; + hipGraph_t graph = nullptr; + hipGraphExec_t graph_exec = nullptr; + std::array node_ops{}; + hipError_t rc = alloc_signal(&ready); + if (rc == hipSuccess) rc = alloc_signal(&done); + if (rc == hipSuccess) rc = hipStreamCreateWithFlags(&stream, hipStreamNonBlocking); + // A runtime that runs the WAIT during capture instead of recording it must not block here. + if (rc == hipSuccess) __atomic_store_n(done, 1ULL, __ATOMIC_RELEASE); + if (rc == hipSuccess) rc = hipStreamBeginCapture(stream, hipStreamCaptureModeRelaxed); + const bool capture_started = rc == hipSuccess; + if (rc == hipSuccess) rc = write_value(stream, ready, 1, explicit_nodes ? &node_ops[0] : nullptr); + if (rc == hipSuccess) rc = wait_done(stream, done, explicit_nodes ? &node_ops[1] : nullptr); + // End the capture even on failure so its stream and graph are not left active during cleanup. + if (capture_started) { + const hipError_t end_rc = hipStreamEndCapture(stream, &graph); + if (rc == hipSuccess) rc = end_rc; + } + if (rc == hipSuccess) rc = hipStreamSynchronize(stream); + size_t nodes = 0; + if (rc == hipSuccess) rc = hipGraphGetNodes(graph, nullptr, &nodes); + if (rc == hipSuccess && nodes != 2) rc = hipErrorNotSupported; + if (rc == hipSuccess) rc = hipGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0); + + bool served = false; + if (rc == hipSuccess) { + __atomic_store_n(ready, 0ULL, __ATOMIC_RELEASE); + __atomic_store_n(done, 0ULL, __ATOMIC_RELEASE); + std::thread responder([&] { + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (__atomic_load_n(ready, __ATOMIC_ACQUIRE) == 0 && std::chrono::steady_clock::now() < deadline) + std::this_thread::yield(); + served = __atomic_load_n(ready, __ATOMIC_ACQUIRE) == 1; + // The probe graph never resets done, so this also releases a launch that arrives after the deadline. + __atomic_store_n(done, 1ULL, __ATOMIC_RELEASE); + }); + rc = hipGraphLaunch(graph_exec, stream); + if (rc == hipSuccess) rc = hipStreamSynchronize(stream); + responder.join(); + } + const bool ok = rc == hipSuccess && served; + + if (graph_exec != nullptr) (void)hipGraphExecDestroy(graph_exec); + if (graph != nullptr) (void)hipGraphDestroy(graph); + if (stream != nullptr) (void)hipStreamDestroy(stream); + free_signal(ready); + free_signal(done); + return ok; +} + +void GpuSync::gpu_submit(uintptr_t stream, uint64_t task) { + auto s = reinterpret_cast(stream); + hipStreamBatchMemOpParams* ops = needs_explicit_nodes(s) ? impl_->submit_ops[task].data() : nullptr; + // Reset done before raising ready, so the coordinator's completion for this step can never be wiped. + hip_check(write_value(s, impl_->done, 0, ops), "write done"); + hip_check(write_value(s, impl_->ready, task, ops ? ops + 1 : nullptr), "write ready"); +} + +void GpuSync::gpu_wait(uintptr_t stream) { + auto s = reinterpret_cast(stream); + hip_check(wait_done(s, impl_->done, needs_explicit_nodes(s) ? &impl_->wait_op : nullptr), "wait done"); +} + +bool GpuSync::host_func_capture_safe() { return false; } + +uint64_t GpuSync::cpu_poll() { + const int64_t task = flag_load_acquire(as_flag(impl_->ready)); + if (task == 0) return 0; + flag_store_release(as_flag(impl_->ready), 0); + return static_cast(task); +} + +void GpuSync::cpu_done() { flag_store_release(as_flag(impl_->done), 1); } + +bool GpuSync::pending() const { + return flag_load_acquire(as_flag(impl_->ready)) != 0 && flag_load_acquire(as_flag(impl_->done)) == 0; +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp new file mode 100644 index 0000000000..273c394609 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp @@ -0,0 +1,159 @@ +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// ------------------------------- dot products ------------------------------- +// dot(weight[bf16], act[bf16], n) -> fp32. The selected impl is a function +// pointer chosen at runtime; the per-row call overhead is negligible vs n. + +using dot_fn = float (*)(const bf16_t*, const bf16_t*, int); + +float dot_scalar(const bf16_t* w, const bf16_t* x, int n) { + float acc = 0.0f; + for (int i = 0; i < n; ++i) acc += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); + return acc; +} + +#if CPU_MOE_X86 +CPU_MOE_TARGET("avx512f") +float dot_avx512f(const bf16_t* w, const bf16_t* x, int n) { + // 4 independent accumulators -> more in-flight loads (memory-level parallelism), + // which is what lifts a bandwidth-bound GEMV toward peak. + __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); + __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); + int i = 0; + for (; i + 64 <= n; i += 64) { + _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); + for (int j = 0; j < 64; j += 16) { + __m256i wi = _mm256_loadu_si256(reinterpret_cast(w + i + j)); + __m256i xi = _mm256_loadu_si256(reinterpret_cast(x + i + j)); + __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16)); + __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16)); + __m512& acc = (j == 0) ? a0 : (j == 16) ? a1 : (j == 32) ? a2 : a3; + acc = _mm512_fmadd_ps(wf, xf, acc); + } + } + for (; i + 16 <= n; i += 16) { + __m256i wi = _mm256_loadu_si256(reinterpret_cast(w + i)); + __m256i xi = _mm256_loadu_si256(reinterpret_cast(x + i)); + __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16)); + __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16)); + a0 = _mm512_fmadd_ps(wf, xf, a0); + } + float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3))); + for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); + return s; +} + +#ifdef CPU_MOE_HAS_AVX512BF16 +CPU_MOE_TARGET("avx512bf16,avx512f") +static inline __m512bh load_bh(const bf16_t* p) { + __m512i raw = _mm512_loadu_si512(reinterpret_cast(p)); + __m512bh out; + std::memcpy(&out, &raw, sizeof(out)); + return out; +} + +CPU_MOE_TARGET("avx512bf16,avx512f") +float dot_avx512bf16(const bf16_t* w, const bf16_t* x, int n) { + // 4 accumulators (128 bf16/iter) for memory-level parallelism + a prefetch nudge. + __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); + __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); + int i = 0; + for (; i + 128 <= n; i += 128) { + _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); + a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i)); + a1 = _mm512_dpbf16_ps(a1, load_bh(w + i + 32), load_bh(x + i + 32)); + a2 = _mm512_dpbf16_ps(a2, load_bh(w + i + 64), load_bh(x + i + 64)); + a3 = _mm512_dpbf16_ps(a3, load_bh(w + i + 96), load_bh(x + i + 96)); + } + for (; i + 32 <= n; i += 32) { + a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i)); + } + float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3))); + for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); + return s; +} +#endif // avx512bf16 available + +// --------- AVX2 + FMA (256-bit fallback for CPUs without AVX-512) ---------- +// Covers Intel 12-14th gen / Arrow Lake (AVX-512 fused off) and AMD Zen<4. The +// bf16->fp32 widen is a zero-extend + <<16; with 4 independent accumulators the +// GEMV is memory-bandwidth bound, same as the AVX-512 path (just half the width). +CPU_MOE_TARGET("avx2,fma") +float dot_avx2(const bf16_t* w, const bf16_t* x, int n) { + __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps(); + __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps(); + int i = 0; + for (; i + 32 <= n; i += 32) { + _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); + for (int j = 0; j < 32; j += 8) { + __m128i wi = _mm_loadu_si128(reinterpret_cast(w + i + j)); + __m128i xi = _mm_loadu_si128(reinterpret_cast(x + i + j)); + __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16)); + __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16)); + __m256& acc = (j == 0) ? a0 : (j == 8) ? a1 : (j == 16) ? a2 : a3; + acc = _mm256_fmadd_ps(wf, xf, acc); + } + } + for (; i + 8 <= n; i += 8) { + __m128i wi = _mm_loadu_si128(reinterpret_cast(w + i)); + __m128i xi = _mm_loadu_si128(reinterpret_cast(x + i)); + __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16)); + __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16)); + a0 = _mm256_fmadd_ps(wf, xf, a0); + } + float s = hsum256(_mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3))); + for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); + return s; +} +#endif // CPU_MOE_X86 + +dot_fn select_dot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 +#ifdef CPU_MOE_HAS_AVX512BF16 + if (t >= ISA_AVX512BF16) return dot_avx512bf16; +#endif + if (t >= ISA_AVX512) return dot_avx512f; + if (t >= ISA_AVX2) return dot_avx2; +#endif + (void)t; + return dot_scalar; +} + +class Bf16Format final : public MoeFormat { + public: + Bf16Format(const MoeParams& p, const FormatArgs& a) + : MoeFormat(p), dot_(select_dot()), gate_up_(a.bank("gate_up")), down_(a.bank("down")) {} + + std::string isa_name() const override { return tier_name(pick_isa()); } + + void gate_up(int layer, int e, int /*tok*/, int i0, int i1, const bf16_t* x, float w_in, + bf16_t* g) const override { + const bf16_t* w = static_cast(tbl_at(gate_up_, layer)) + (size_t)e * (2 * p.I) * p.H; + gated_rows(p, i0, i1, w_in, g, [&](int row) { return dot_(w + (size_t)row * p.H, x, p.H); }); + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const bf16_t* w = static_cast(tbl_at(down_, layer)); + for (int h = h0; h < h1; ++h) + for (int j = 0; j < n; ++j) + out[j * kRowBlock + h - h0] = dot_(w + ((size_t)routes[j].e * p.H + h) * p.I, routes[j].g, p.I); + } + + private: + dot_fn dot_; + // bf16 rows: gate_up [E, 2I, H], down [E, H, I]. + const uint64_t* gate_up_; + const uint64_t* down_; +}; + +} // namespace + +std::unique_ptr make_bf16_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp new file mode 100644 index 0000000000..a7afe9804d --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp @@ -0,0 +1,269 @@ +#include + +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// ----------------------- DeepSeek-V4 ds_fp4 (W4A8) --------------------------- +// Row-major e2m1 (2/byte, low nibble first) + e8m0 per-32 block scale, no global +// (w = E2M1[code] * 2^(e8m0-127)); activations are FP8-e4m3 round-tripped (per-128 +// block, ue8m0 scale) before each GEMM. Matches kernel/triton/dsv4 (fused_moe + +// fp8_linear): silu(clamp(gate,max=lim)) * clamp(up,-lim,lim), router weight on the +// down output. + +// Activations are pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]) once per +// token/route and reused across all output rows. This drops the hot dot to one +// vpmovzxbd + two vpermps per 32 weights (1.5 shuffle ops / 16 vs 4 for a bf16, +// dup-permute, per-element gather), so the row-major fp4 GEMV stops being port-5 +// bound and approaches the bf16 memory-bandwidth ceiling. +using dsdot_fn = float (*)(const uint8_t*, const uint8_t*, const float*, const float*, int, + const float*, const float*); + +float dot_dsfp4_scalar(const uint8_t* packed, const uint8_t* scale, const float* xe, + const float* xo, int K, const float* e2m1, const float* e8m0) { + float acc = 0.0f; + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const float sc = e8m0[scale[b]]; + const uint8_t* pk = packed + (size_t)b * 16; // 16 bytes = 32 codes + const float* xeb = xe + (size_t)b * 16; + const float* xob = xo + (size_t)b * 16; + float bsum = 0.0f; + for (int j = 0; j < 16; ++j) { + const uint8_t byte = pk[j]; + bsum += e2m1[byte & 0xF] * xeb[j]; // low nibble -> even-K activation + bsum += e2m1[byte >> 4] * xob[j]; // high nibble -> odd-K activation + } + acc += sc * bsum; + } + return acc; +} + +#if CPU_MOE_X86 +// One 32-block: 16 bytes -> 16 low + 16 high nibble values via two vpermps, times +// the pre-split even/odd fp32 activations, folded by the per-32 e8m0 scale. +CPU_MOE_TARGET("avx512f") +inline __m512 dsfp4_blk(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut, + __m512i loma, float sc) { + __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast(pk))); + __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); + __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); + __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob))); + return _mm512_mul_ps(prod, _mm512_set1_ps(sc)); +} + +CPU_MOE_TARGET("avx512f") +float dot_dsfp4_avx512(const uint8_t* packed, const uint8_t* scale, const float* xe, + const float* xo, int K, const float* e2m1, const float* e8m0) { + const __m512 lut = _mm512_loadu_ps(e2m1); + const __m512i loma = _mm512_set1_epi32(0xF); + __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); + const int nb = K / 32; // 16 packed bytes + one e8m0 scale per block + int b = 0; + for (; b + 2 <= nb; b += 2) { // two independent accumulators hide FMA latency + acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16, + xo + (size_t)b * 16, lut, loma, e8m0[scale[b]])); + acc1 = _mm512_add_ps(acc1, dsfp4_blk(packed + (size_t)(b + 1) * 16, xe + (size_t)(b + 1) * 16, + xo + (size_t)(b + 1) * 16, lut, loma, e8m0[scale[b + 1]])); + } + for (; b < nb; ++b) + acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16, + xo + (size_t)b * 16, lut, loma, e8m0[scale[b]])); + return _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1)); +} + +// AVX2: a 32-K block is 16 bytes -> two 8-lane halves (8 even + 8 odd each). +CPU_MOE_TARGET("avx2,fma") +inline __m256 dsfp4_half_avx2(const uint8_t* pk, const float* xeb, const float* xob, __m256 mag8) { + __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); + __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); + __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); + return _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); +} + +CPU_MOE_TARGET("avx2,fma") +float dot_dsfp4_avx2(const uint8_t* packed, const uint8_t* scale, const float* xe, + const float* xo, int K, const float* e2m1, const float* e8m0) { + const __m256 mag8 = _mm256_loadu_ps(e2m1); + __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* pk = packed + (size_t)b * 16; + const float* xeb = xe + (size_t)b * 16; + const float* xob = xo + (size_t)b * 16; + const __m256 sc = _mm256_set1_ps(e8m0[scale[b]]); + acc0 = _mm256_fmadd_ps(dsfp4_half_avx2(pk, xeb, xob, mag8), sc, acc0); + acc1 = _mm256_fmadd_ps(dsfp4_half_avx2(pk + 8, xeb + 8, xob + 8, mag8), sc, acc1); + } + return hsum256(_mm256_add_ps(acc0, acc1)); +} +#endif + +dsdot_fn select_dsdot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (t >= ISA_AVX512) return dot_dsfp4_avx512; + if (t >= ISA_AVX2) return dot_dsfp4_avx2; +#endif + (void)t; + return dot_dsfp4_scalar; +} + +// Round a clamped |x|<=448 to nearest float8-e4m3 (RNE), back to fp32. Matches +// torch.float8_e4m3fn / triton .to(float8e4nv). +inline float e4m3_round(float x) { + const float sign = x < 0.0f ? -1.0f : 1.0f; + const float a = std::fabs(x); + if (a == 0.0f) return 0.0f; + if (a >= 448.0f) return sign * 448.0f; + int e; + std::frexp(a, &e); // a in [2^(e-1), 2^e) + float step = std::ldexp(1.0f, e - 4); + const float min_step = std::ldexp(1.0f, -9); // e4m3 subnormal step (2^-9) + if (step < min_step) step = min_step; + float r = std::nearbyint(a / step) * step; + if (r > 448.0f) r = 448.0f; + return sign * r; +} + +// IEEE ceil(log2(v)) for v>0 (matches dsv4 _log2_ceil / fast_round_scale). +inline int ceil_log2_pos(float v) { + uint32_t bits; + std::memcpy(&bits, &v, sizeof(bits)); + const int exp = (int)((bits >> 23) & 0xFF); + const int man = (int)(bits & 0x7FFFFF); + return exp - 127 + (man != 0 ? 1 : 0); +} + +// DeepSeek-V4 / V4.1 activation FP8 round-trip (bf16 in/out): per `block` (the checkpoint's +// fp8 block: 128 on V4, 32 on V4.1), s = 2^ceil(log2(max(|x|,1e-4)/448)); +// y = round_e4m3(clamp(x/s,+-448)) * s. +void fp8_roundtrip_bf16(const bf16_t* src, bf16_t* dst, int K, int block) { + for (int b0 = 0; b0 < K; b0 += block) { + const int b1 = std::min(K, b0 + block); + float amax = 1e-4f; + for (int i = b0; i < b1; ++i) amax = std::max(amax, std::fabs(bf16_to_f32(src[i]))); + const float s = std::ldexp(1.0f, ceil_log2_pos(amax * (1.0f / 448.0f))); + const float inv_s = 1.0f / s; + for (int i = b0; i < b1; ++i) { + float q = bf16_to_f32(src[i]) * inv_s; + q = std::min(448.0f, std::max(-448.0f, q)); + dst[i] = f32_to_bf16(e4m3_round(q) * s); + } + } +} + +class DsFp4Format final : public MoeFormat { + public: + DsFp4Format(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), dot_(select_dsdot()) { + gate_up_ = a.bank("gate_up"); + gu_scale_ = a.bank("gate_up_scale"); + down_ = a.bank("down"); + dn_scale_ = a.bank("down_scale"); + act_block_ = static_cast(a.option("act_block")); + input_prequant_ = a.option("input_prequant", 0) != 0; + const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k; + xq_.resize(tokens * p.H); + xe_.resize(tokens * (p.H / 2)); + xo_.resize(tokens * (p.H / 2)); + ge_.resize(routes * (p.I / 2)); + go_.resize(routes * (p.I / 2)); + for (int i = 0; i < 16; ++i) e2m1_lut_[i] = kE2M1[i]; + for (int i = 0; i < 256; ++i) e8m0_lut_[i] = std::ldexp(1.0f, std::min(i, 254) - 127); + } + + std::string isa_name() const override { return tier_name(pick_isa()); } + + + // DSV4 FP8-round-trips the input before the gate_up GEMV, unless the GPU already did. + void prep_input(int tok, const bf16_t* x) override { + const bf16_t* src = x; + if (!input_prequant_) { + bf16_t* xq = xq_.data() + (size_t)tok * p.H; + fp8_roundtrip_bf16(x, xq, p.H, act_block_); + src = xq; + } + deinterleave_bf16_f32(src, xe_.data() + (size_t)tok * (p.H / 2), xo_.data() + (size_t)tok * (p.H / 2), p.H); + } + + bool preps_inter() const override { return true; } + + void prep_inter(int64_t r, bf16_t* g) override { + fp8_roundtrip_bf16(g, g, p.I, act_block_); + deinterleave_bf16_f32(g, ge_.data() + (size_t)r * (p.I / 2), go_.data() + (size_t)r * (p.I / 2), p.I); + } + + void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float /*w_in*/, + bf16_t* g) const override { + const int H = p.H, I = p.I; + const uint8_t* gu_packed_l = static_cast(tbl_at(gate_up_, layer)); + const uint8_t* gu_scale_l = static_cast(tbl_at(gu_scale_, layer)); + const int N2 = 2 * I, Hh = H / 2, Hs = H / 32; + const float* xe = xe_.data() + (size_t)tok * (H / 2); + const float* xo = xo_.data() + (size_t)tok * (H / 2); + 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 float lim = p.swiglu_limit; + for (int i = i0; i < i1; ++i) { + // gate_up is stored bf16 by the reference GEMV before swiglu; round to match. + float gate = bf16_to_f32(f32_to_bf16( + dot_(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(dot_( + 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)); // silu(gate) + g[i] = f32_to_bf16(glu * up); + } + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const int H = p.H, I = p.I, Ih = I / 2, Is = I / 32; + const uint8_t* dn_packed_l = static_cast(tbl_at(down_, layer)); + const uint8_t* dn_scale_l = static_cast(tbl_at(dn_scale_, layer)); + for (int h = h0; h < h1; ++h) { + for (int j = 0; j < n; ++j) { + const size_t e = (size_t)routes[j].e, r = (size_t)routes[j].r; + const uint8_t* dp = dn_packed_l + e * (size_t)H * Ih + (size_t)h * Ih; + const uint8_t* ds = dn_scale_l + e * (size_t)H * Is + (size_t)h * Is; + out[j * kRowBlock + h - h0] = dot_(dp, ds, ge_.data() + r * Ih, go_.data() + r * Ih, I, e2m1_lut_, e8m0_lut_); + } + } + } + + // The reference rounds each route's weighted down output to bf16 before the fp32 sum over routes. + bool rounds_routes() const override { return true; } + + private: + dsdot_fn dot_; + // e2m1 rows [E,2I,H/2] with e8m0 scales per 32 along K [E,2I,H/32]; down likewise with K = I. + const uint64_t* gate_up_ = nullptr; + const uint64_t* gu_scale_ = nullptr; + const uint64_t* down_ = nullptr; + const uint64_t* dn_scale_ = nullptr; + // Activation FP8 round-trip block, the checkpoint's fp8 block. + int act_block_ = 128; + // The caller already FP8-round-tripped the input activations on the GPU + // (same reference grid), so prep_input must not repeat it on the host-callback + // thread. That scalar per-element pass is single-threaded ON THE DECODE CRITICAL + // PATH (~0.3ms/layer at H=4096, every worker and the GPU waiting on it); moving + // it to a captured GPU elementwise kernel removes it while keeping the official + // W4A8 numerics bit-exact (see cpu_executor.py). + bool input_prequant_ = false; + float e2m1_lut_[16]; + float e8m0_lut_[256]; + AlignedVector xq_; + AlignedVector xe_, xo_, ge_, go_; +}; + +} // namespace + +std::unique_ptr make_ds_fp4_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp new file mode 100644 index 0000000000..d27fb1868e --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp @@ -0,0 +1,203 @@ +#include +#include + +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// ------------------------ block-FP8 (W8A16) dequant ------------------------ +// fp8-e4m3 weights row-major [rows, K] plus one bf16 scale per 128x128 block. The scale tensor is named weight_scale_inv but multiplies, as in kernel/triton/fp8_blockscale_moe.py. +// Whole 128-blocks read the activations as fp32 even/odd halves (xe[m] = x[2m], xo[m] = x[2m+1]); a ragged tail reads the bf16 row x. +// 0x7F/0xFF (NaN in torch float8_e4m3fn) decode to +-480 in every kernel, as e4m3_decode does. +using fp8_block_dot_fn = float (*)(const uint8_t*, const float*, const float*, const bf16_t*, + const bf16_t*, int, const float*); + +constexpr int FP8_BLK = 128; // K elements covered by one weight scale + +inline float fp8_block_tail(const uint8_t* w, const bf16_t* x, int k0, int K, const float* e4m3) { + float blk = 0.0f; + for (int k = k0; k < K; ++k) blk += e4m3[w[k]] * bf16_to_f32(x[k]); + return blk; +} + +float dot_fp8_block_scalar(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, + const bf16_t* s, int K, const float* e4m3) { + float acc = 0.0f; + int b = 0, k0 = 0; + for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { + // 8 independent sums: a single running sum is bound by the add latency + float a[8] = {}; + const uint8_t* wb = w + k0; + const float* pe = xe + k0 / 2; + const float* po = xo + k0 / 2; + for (int j = 0; j < FP8_BLK / 2; j += 4) + for (int l = 0; l < 4; ++l) { + a[2 * l] += e4m3[wb[2 * (j + l)]] * pe[j + l]; + a[2 * l + 1] += e4m3[wb[2 * (j + l) + 1]] * po[j + l]; + } + acc += (((a[0] + a[1]) + (a[2] + a[3])) + ((a[4] + a[5]) + (a[6] + a[7]))) * + bf16_to_f32(s[b]); + } + if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); + return acc; +} + +#if CPU_MOE_X86 +// A byte b read as fp16 bits (b & 0x80) << 8 | (b & 0x7F) << 7 is its e4m3 value times 2^-8, subnormals included, so the block scale carries the 256. +// In a 16-bit lane holding bytes (2m, 2m+1), srai 1 then & 0xBF80 builds that for byte 2m+1; a left shift by 8 first does byte 2m. Every AVX2 CPU has F16C. +CPU_MOE_TARGET("avx2,fma,f16c") +float dot_fp8_block_avx2(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, + const bf16_t* s, int K, const float* e4m3) { + const __m256i keep = _mm256_set1_epi16((short)0xBF80); + __m256 total = _mm256_setzero_ps(); + int b = 0, k0 = 0; + for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { + _mm_prefetch(reinterpret_cast(w + k0) + PF_AHEAD, _MM_HINT_T0); + __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps(); + __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps(); + for (int j = 0; j < FP8_BLK; j += 32) { + const __m256i v = _mm256_loadu_si256(reinterpret_cast(w + k0 + j)); + const __m256i odd = _mm256_and_si256(_mm256_srai_epi16(v, 1), keep); + const __m256i even = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v, 8), 1), keep); + const float* pe = xe + (k0 + j) / 2; + const float* po = xo + (k0 + j) / 2; + a0 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(even)), _mm256_loadu_ps(pe), a0); + a1 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(even, 1)), _mm256_loadu_ps(pe + 8), a1); + a2 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(odd)), _mm256_loadu_ps(po), a2); + a3 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(odd, 1)), _mm256_loadu_ps(po + 8), a3); + } + const __m256 blk = _mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3)); + total = _mm256_fmadd_ps(blk, _mm256_set1_ps(bf16_to_f32(s[b]) * 256.0f), total); + } + float acc = hsum256(total); + if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); + return acc; +} + +CPU_MOE_TARGET("avx512f,avx2,fma,f16c") +float dot_fp8_block_avx512f(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, + const bf16_t* s, int K, const float* e4m3) { + const __m256i keep = _mm256_set1_epi16((short)0xBF80); + __m512 total = _mm512_setzero_ps(); + int b = 0, k0 = 0; + for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { + _mm_prefetch(reinterpret_cast(w + k0) + PF_AHEAD, _MM_HINT_T0); + __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); + __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); + for (int j = 0; j < FP8_BLK; j += 64) { + const __m256i v0 = _mm256_loadu_si256(reinterpret_cast(w + k0 + j)); + const __m256i v1 = _mm256_loadu_si256(reinterpret_cast(w + k0 + j + 32)); + const __m256i o0 = _mm256_and_si256(_mm256_srai_epi16(v0, 1), keep); + const __m256i e0 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v0, 8), 1), keep); + const __m256i o1 = _mm256_and_si256(_mm256_srai_epi16(v1, 1), keep); + const __m256i e1 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v1, 8), 1), keep); + const float* pe = xe + (k0 + j) / 2; + const float* po = xo + (k0 + j) / 2; + a0 = _mm512_fmadd_ps(_mm512_cvtph_ps(e0), _mm512_loadu_ps(pe), a0); + a1 = _mm512_fmadd_ps(_mm512_cvtph_ps(o0), _mm512_loadu_ps(po), a1); + a2 = _mm512_fmadd_ps(_mm512_cvtph_ps(e1), _mm512_loadu_ps(pe + 16), a2); + a3 = _mm512_fmadd_ps(_mm512_cvtph_ps(o1), _mm512_loadu_ps(po + 16), a3); + } + const __m512 blk = _mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3)); + total = _mm512_fmadd_ps(blk, _mm512_set1_ps(bf16_to_f32(s[b]) * 256.0f), total); + } + float acc = _mm512_reduce_add_ps(total); + if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); + return acc; +} +#endif // CPU_MOE_X86 + +fp8_block_dot_fn select_fp8_block_dot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + // the avx512bf16 tier keeps the fp32 kernel: on Sapphire Rapids it is faster than widening to bf16 for dpbf16 + if (t >= ISA_AVX512) return dot_fp8_block_avx512f; + if (t >= ISA_AVX2) return dot_fp8_block_avx2; +#endif + (void)t; + return dot_fp8_block_scalar; +} + +class Fp8BlockFormat final : public MoeFormat { + public: + Fp8BlockFormat(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), dot_(select_fp8_block_dot()) { + gate_up_ = a.bank("gate_up"); + gu_scale_ = a.bank("gate_up_scale"); + down_ = a.bank("down"); + dn_scale_ = a.bank("down_scale"); + gu_scale_stride_ = a.option("gu_scale_stride"); + dn_scale_stride_ = a.option("dn_scale_stride"); + const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k; + xe_.resize(tokens * (p.H / 2)); + xo_.resize(tokens * (p.H / 2)); + ge_.resize(routes * (p.I / 2)); + go_.resize(routes * (p.I / 2)); + for (int i = 0; i < 256; ++i) e4m3_lut_[i] = e4m3_decode((uint8_t)i); + } + + std::string isa_name() const override { return tier_name(pick_isa()); } + + + void prep_input(int tok, const bf16_t* x) override { + deinterleave_bf16_f32(x, xe_.data() + (size_t)tok * (p.H / 2), xo_.data() + (size_t)tok * (p.H / 2), p.H); + } + + bool preps_inter() const override { return true; } + + void prep_inter(int64_t r, bf16_t* g) override { + deinterleave_bf16_f32(g, ge_.data() + (size_t)r * (p.I / 2), go_.data() + (size_t)r * (p.I / 2), p.I); + } + + // The halves cover whole 128-blocks; the bf16 row covers a ragged tail. + void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* x, float w_in, + bf16_t* g) const override { + const int H = p.H, I = p.I; + const uint8_t* packed = static_cast(tbl_at(gate_up_, layer)); + const bf16_t* scale = static_cast(tbl_at(gu_scale_, layer)); + const float* xe = xe_.data() + (size_t)tok * (H / 2); + const float* xo = xo_.data() + (size_t)tok * (H / 2); + gated_rows(p, i0, i1, w_in, g, [&](int row) { + const uint8_t* w = packed + ((size_t)e * (2 * I) + row) * (size_t)H; + const bf16_t* sc = scale + ((size_t)e * ((2 * I + FP8_BLK - 1) / FP8_BLK) + row / FP8_BLK) * + (size_t)gu_scale_stride_; + return dot_(w, xe, xo, x, sc, H, e4m3_lut_); + }); + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const int H = p.H, I = p.I; + const uint8_t* packed = static_cast(tbl_at(down_, layer)); + const bf16_t* scale = static_cast(tbl_at(dn_scale_, layer)); + for (int h = h0; h < h1; ++h) { + for (int j = 0; j < n; ++j) { + const size_t e = (size_t)routes[j].e, r = (size_t)routes[j].r; + const uint8_t* w = packed + (e * H + h) * (size_t)I; + const bf16_t* sc = scale + (e * ((H + FP8_BLK - 1) / FP8_BLK) + h / FP8_BLK) * + (size_t)dn_scale_stride_; + out[j * kRowBlock + h - h0] = + dot_(w, ge_.data() + r * (I / 2), go_.data() + r * (I / 2), routes[j].g, sc, I, e4m3_lut_); + } + } + } + + private: + fp8_block_dot_fn dot_; + // e4m3 rows [E,2I,H] and [E,H,I] with one bf16 scale per 128x128 block. + const uint64_t* gate_up_ = nullptr; + const uint64_t* gu_scale_ = nullptr; + const uint64_t* down_ = nullptr; + const uint64_t* dn_scale_ = nullptr; + // Row strides of the scale tensors, including any padding of the host bank layout. + int64_t gu_scale_stride_ = 0, dn_scale_stride_ = 0; + float e4m3_lut_[256]; + AlignedVector xe_, xo_, ge_, go_; +}; + +} // namespace + +std::unique_ptr make_fp8_block_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp new file mode 100644 index 0000000000..484d6dc001 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp @@ -0,0 +1,247 @@ +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// ------------------------- mxfp4 (gpt-oss) GEMV ----------------------------- +// Transposed split-K layout: blk[Kpairs, N2] (N innermost), scl[Kpairs/16, N2] +// e8m0 per 32-K. Computes out[c] = sum_kb (E2M1[lo]*x[2kb] + E2M1[hi]*x[2kb+1]) +// * 2^(e8m0-127) for a contiguous column tile (blk/scl already offset to col 0 of +// the tile). Vectorized over N (16 columns / __m512), K stays the outer (cache- +// sequential) loop. Used by both gate_up (K=H) and down (K=I). +using mxgemv_fn = void (*)(float*, const uint8_t*, const uint8_t*, const bf16_t*, int, int, + int, const float*, const float*); + +void mxfp4_gemv_scalar(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + for (int c = 0; c < ncol; ++c) out[c] = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t* w = blk + (size_t)kb * N2; + const uint8_t* s = scl + (size_t)(kb >> 4) * N2; + const float xl = bf16_to_f32(x[2 * kb]); + const float xh = bf16_to_f32(x[2 * kb + 1]); + for (int c = 0; c < ncol; ++c) { + const uint8_t byte = w[c]; + out[c] += (e2m1[byte & 0xF] * xl + e2m1[byte >> 4] * xh) * e8m0[s[c]]; + } + } +} + +#if CPU_MOE_X86 +CPU_MOE_TARGET("avx512f") +void mxfp4_gemv_avx512(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + (void)e8m0; // e8m0[c]=2^(c-127) computed via bit construction (no gather) + const __m512 lut = _mm512_loadu_ps(e2m1); + const __m512i loma = _mm512_set1_epi32(0xF); + // K-outer / N-inner: each kb cache line is read once and all live column chunks + // (up to 4 -> 64 cols) accumulate from registers, so DRAM/L2 stream the tile once. + int c0 = 0; + for (; c0 + 16 <= ncol; c0 += 64) { + const int nchunk = std::min(4, (ncol - c0) / 16); + __m512 acc[4]; + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_setzero_ps(); + for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = 32 K = one scale row + __m512 sc[4]; + for (int ci = 0; ci < nchunk; ++ci) { + __m128i sraw = _mm_loadu_si128(reinterpret_cast( + scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 16)); + sc[ci] = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu8_epi32(sraw), 23)); + } + __m512 blk_acc[4]; + for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm512_setzero_ps(); + for (int kk = 0; kk < 16; ++kk) { + const int kb = kblk + kk; + const uint8_t* wbase = blk + (size_t)kb * N2 + c0; + // The transposed layout strides K by N2 bytes; prefetch ahead so the strided + // reads are not exposed to DRAM latency (the HW streamer misses big strides). + constexpr int PFD = 8; + if (kb + PFD < Kpairs) + _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), + _MM_HINT_T0); + const __m512 xl = _mm512_set1_ps(bf16_to_f32(x[2 * kb])); + const __m512 xh = _mm512_set1_ps(bf16_to_f32(x[2 * kb + 1])); + for (int ci = 0; ci < nchunk; ++ci) { + __m512i wi = _mm512_cvtepu8_epi32( + _mm_loadu_si128(reinterpret_cast(wbase + ci * 16))); + __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); + __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); + blk_acc[ci] = _mm512_fmadd_ps(vlo, xl, blk_acc[ci]); + blk_acc[ci] = _mm512_fmadd_ps(vhi, xh, blk_acc[ci]); + } + } + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); + } + for (int ci = 0; ci < nchunk; ++ci) _mm512_storeu_ps(out + c0 + ci * 16, acc[ci]); + } + for (int c = c0; c < ncol; ++c) { // tail columns (< 16) + float o = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t byte = blk[(size_t)kb * N2 + c]; + uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; + float sc; + std::memcpy(&sc, &bits, 4); + o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + + e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; + } + out[c] = o; + } +} + +CPU_MOE_TARGET("avx2,fma") +void mxfp4_gemv_avx2(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + (void)e8m0; // e8m0[s]=2^(s-127) built via s<<23 (no gather) + const __m256 mag8 = _mm256_loadu_ps(e2m1); + int c0 = 0; + for (; c0 + 8 <= ncol; c0 += 32) { // up to 4 chunks of 8 = 32 cols + const int nchunk = std::min(4, (ncol - c0) / 8); + __m256 acc[4]; + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_setzero_ps(); + for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = one scale row + __m256 sc[4]; + for (int ci = 0; ci < nchunk; ++ci) { + __m128i sraw = _mm_loadl_epi64(reinterpret_cast( + scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 8)); + sc[ci] = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu8_epi32(sraw), 23)); + } + __m256 blk_acc[4]; + for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm256_setzero_ps(); + for (int kk = 0; kk < 16; ++kk) { + const int kb = kblk + kk; + const uint8_t* wbase = blk + (size_t)kb * N2 + c0; + constexpr int PFD = 8; + if (kb + PFD < Kpairs) + _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), + _MM_HINT_T0); + const __m256 xl = _mm256_set1_ps(bf16_to_f32(x[2 * kb])); + const __m256 xh = _mm256_set1_ps(bf16_to_f32(x[2 * kb + 1])); + for (int ci = 0; ci < nchunk; ++ci) { + __m256i wi = _mm256_cvtepu8_epi32( + _mm_loadl_epi64(reinterpret_cast(wbase + ci * 8))); + __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); + __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); + blk_acc[ci] = _mm256_fmadd_ps(vlo, xl, blk_acc[ci]); + blk_acc[ci] = _mm256_fmadd_ps(vhi, xh, blk_acc[ci]); + } + } + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); + } + for (int ci = 0; ci < nchunk; ++ci) _mm256_storeu_ps(out + c0 + ci * 8, acc[ci]); + } + for (int c = c0; c < ncol; ++c) { // tail columns (< 8); none when ncol%8==0 + float o = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t byte = blk[(size_t)kb * N2 + c]; + uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; + float sc; + std::memcpy(&sc, &bits, 4); + o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + + e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; + } + out[c] = o; + } +} +#endif + +mxgemv_fn select_mxgemv() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (t >= ISA_AVX512) return mxfp4_gemv_avx512; + if (t >= ISA_AVX2) return mxfp4_gemv_avx2; +#endif + (void)t; + return mxfp4_gemv_scalar; +} + +// ----------------------------- mxfp4 (gpt-oss) ----------------------------- +// Transposed split-K layout (N innermost), so the GEMV streams K and accumulates +// a contiguous block of N output columns -> cache-line-efficient, no repack, no +// extra host memory. Pass 1 fuses gate_up + clamped-swiglu(+bias); pass 2 adds the +// down bias, and the executor applies the router weight and sums the routes. +// +// 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. +class Mxfp4Format final : public MoeFormat { + public: + Mxfp4Format(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), gemv_(select_mxgemv()) { + gate_up_ = a.bank("gate_up"); + gu_scale_ = a.bank("gate_up_scale"); + gu_bias_ = a.bank("gate_up_bias"); + down_ = a.bank("down"); + dn_scale_ = a.bank("down_scale"); + dn_bias_ = a.bank("down_bias"); + for (int i = 0; i < 16; ++i) e2m1_lut_[i] = kE2M1[i]; + // e8m0 (mxfp4 block scale) = 2^(s-127); the GPU GEMV clamps s to [0,254]. + for (int i = 0; i < 256; ++i) e8m0_lut_[i] = std::ldexp(1.0f, std::min(i, 254) - 127); + } + + std::string isa_name() const override { return tier_name(pick_isa()); } + + void gate_up(int layer, int e, int /*tok*/, int i0, int i1, const bf16_t* x, float /*w_in*/, + bf16_t* g) const override { + const int H = p.H, I = p.I; + const uint8_t* gu_packed_l = static_cast(tbl_at(gate_up_, layer)); + const uint8_t* gu_scale_l = static_cast(tbl_at(gu_scale_, layer)); + const bf16_t* gu_bias_l = static_cast(tbl_at(gu_bias_, layer)); + const int N2 = 2 * I; // gate_up output width (gate/up interleaved) + const int Hh = H / 2; // packed-K rows + const int nunit = i1 - i0; // intermediate units owned by this tile + const int col0 = 2 * i0; // first gate_up column + const int ncol = 2 * nunit; // gate_up columns owned by this tile + 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; + float gu[2 * kRowBlock]; + gemv_(gu, blk_e + col0, scl_e + col0, x, Hh, N2, ncol, e2m1_lut_, e8m0_lut_); + const bf16_t* bias_e = gu_bias_l + (size_t)e * N2 + col0; + const float lim = p.swiglu_limit, alpha = p.swiglu_alpha; + 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)); // gate * sigmoid(alpha*gate) + g[i0 + j] = f32_to_bf16(glu * (up + 1.0f)); + } + } + + // The transposed layout already streams a whole column block per call, so the routes go one after another. + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const int H = p.H, I = p.I; + const int Ih = I / 2; // packed-K rows + const int nh = h1 - h0; + const uint8_t* dn_packed_l = static_cast(tbl_at(down_, layer)); + const uint8_t* dn_scale_l = static_cast(tbl_at(dn_scale_, layer)); + const bf16_t* dn_bias_l = static_cast(tbl_at(dn_bias_, layer)); + for (int j = 0; j < n; ++j) { + const size_t e = (size_t)routes[j].e; + const uint8_t* blk_e = dn_packed_l + e * Ih * H; + const uint8_t* scl_e = dn_scale_l + e * (size_t)(I / 32) * H; + float* part = out + j * kRowBlock; + gemv_(part, blk_e + h0, scl_e + h0, routes[j].g, Ih, H, nh, e2m1_lut_, e8m0_lut_); + const bf16_t* bias_e = dn_bias_l + e * H + h0; + for (int c = 0; c < nh; ++c) part[c] += bf16_to_f32(bias_e[c]); + } + } + + private: + mxgemv_fn gemv_; + // gpt-oss split-K layout, N innermost: blocks [E,H/2,2I], e8m0 scales [E,H/32,2I], bf16 bias [E,2I]; down likewise with K = I, N = H. + const uint64_t* gate_up_ = nullptr; + const uint64_t* gu_scale_ = nullptr; + const uint64_t* gu_bias_ = nullptr; + const uint64_t* down_ = nullptr; + const uint64_t* dn_scale_ = nullptr; + const uint64_t* dn_bias_ = nullptr; + float e2m1_lut_[16]; + float e8m0_lut_[256]; +}; + +} // namespace + +std::unique_ptr make_mxfp4_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp new file mode 100644 index 0000000000..429955faf3 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp @@ -0,0 +1,466 @@ +#include + +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// --------------------------- NVFP4 (W4A16) dequant --------------------------- +// Weights: e2m1 4-bit codes (2/byte, low nibble first), per-16 block scale in +// fp8-e4m3, per-output-row global scale in fp16. Dequant matches the GPU kernels +// (freetoken/kernel/triton/nvfp4_dequant.py): w = E2M1[code] * e4m3(scale) * global. +// Activations stay bf16 (W4A16); the GEMV dequantizes weights inside the K-loop. + + +// e2m1 * 2 as exact int8 (all e2m1 values are multiples of 0.5). Used by the AVX-VNNI +// W4A8 path: nibble -> int8 weight via PSHUFB LUT, then VPDPBUSD against int8 activations; +// the *2 is undone by a 0.5 folded into the final scale. Mirrors ggml's kvalues_mxfp4. +alignas(16) const int8_t kE2M1x2[16] = {0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12}; + +// Activations pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]); see the +// ds_fp4 dot in ds_fp4.cpp for why (drops the hot loop to ~1.5 shuffle ops / 16 weights). +using nvdot_fn = float (*)(const uint8_t*, const uint8_t*, float, const float*, const float*, + int, const float*, const float*); + +float dot_nvfp4_scalar(const uint8_t* packed, const uint8_t* scale, float global, + const float* xe, const float* xo, int K, const float* e2m1, + const float* e4m3) { + float acc = 0.0f; + const int nb = K / 16; + for (int b = 0; b < nb; ++b) { + const float bs = e4m3[scale[b]]; + const uint8_t* pk = packed + (size_t)b * 8; + const float* xeb = xe + (size_t)b * 8; // 16 K -> 8 even + 8 odd + const float* xob = xo + (size_t)b * 8; + float bsum = 0.0f; + for (int j = 0; j < 8; ++j) { + const uint8_t byte = pk[j]; + bsum += e2m1[byte & 0xF] * xeb[j]; + bsum += e2m1[byte >> 4] * xob[j]; + } + acc += bs * bsum; + } + return acc * global; +} + +// ---- NVFP4 W4A8 (int8 activations) dot: nibble->int8 via LUT, per-16 act scale ---- +// asi8: int8 activations laid out per-16 block as [even(8), odd(8)]; asb[b] = per-block +// activation scale (absmax/127). Result folds the e2m1*2 -> *0.5 into the scale. +using nvi8dot_fn = float (*)(const uint8_t*, const uint8_t*, float, const int8_t*, int, + const float*, const float*); + +[[maybe_unused]] float dot_nvfp4_i8_scalar(const uint8_t* packed, const uint8_t* scale, + float global, const int8_t* asi8, int K, const float* e4m3, + const float* asb) { + float acc = 0.0f; + const int nb = K / 16; + for (int b = 0; b < nb; ++b) { + const uint8_t* pk = packed + (size_t)b * 8; + const int8_t* ae = asi8 + (size_t)b * 16; // even(8) + const int8_t* ao = ae + 8; // odd(8) + int isum = 0; + for (int j = 0; j < 8; ++j) { + isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j]; + isum += (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; + } + acc += (e4m3[scale[b]] * asb[b]) * (float)isum; + } + return acc * (0.5f * global); +} + +#if CPU_MOE_X86 +// Two 16-K blocks (16 packed bytes) per iter: lo nibbles -> even-K, hi -> odd-K, +// gathered via two vpermps. The per-16 e4m3 scale differs across the two blocks, so +// it is applied per lane (low 8 lanes = block b, high 8 = block b+1). +CPU_MOE_TARGET("avx512f") +inline __m512 nvfp4_blk2(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut, + __m512i loma, float s0, float s1) { + __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast(pk))); + __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); + __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); + __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob))); + // lanes 0-7 (block b) -> s0, lanes 8-15 (block b+1) -> s1 (pure AVX512F mask move) + __m512 scv = _mm512_mask_mov_ps(_mm512_set1_ps(s0), 0xFF00, _mm512_set1_ps(s1)); + return _mm512_mul_ps(prod, scv); +} + +CPU_MOE_TARGET("avx512f") +float dot_nvfp4_avx512(const uint8_t* packed, const uint8_t* scale, float global, + const float* xe, const float* xo, int K, const float* e2m1, + const float* e4m3) { + const __m512 lut = _mm512_loadu_ps(e2m1); // 16 e2m1 values + const __m512i loma = _mm512_set1_epi32(0xF); + __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); + const int nb = K / 16; // 8 packed bytes + one e4m3 scale per block + int b = 0; + for (; b + 4 <= nb; b += 4) { // two blk2 calls -> 4 blocks, 2 accumulators + acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8, + xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]])); + acc1 = _mm512_add_ps(acc1, nvfp4_blk2(packed + (size_t)(b + 2) * 8, xe + (size_t)(b + 2) * 8, + xo + (size_t)(b + 2) * 8, lut, loma, e4m3[scale[b + 2]], e4m3[scale[b + 3]])); + } + for (; b + 2 <= nb; b += 2) + acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8, + xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]])); + float s = _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1)); + for (; b < nb; ++b) { // odd final 16-K block + const uint8_t* pk = packed + (size_t)b * 8; + const float* xeb = xe + (size_t)b * 8; + const float* xob = xo + (size_t)b * 8; + float bsum = 0.0f; + for (int j = 0; j < 8; ++j) { + bsum += e2m1[pk[j] & 0xF] * xeb[j] + e2m1[pk[j] >> 4] * xob[j]; + } + s += e4m3[scale[b]] * bsum; + } + return s * global; +} + +// AVX2: one 16-K block (8 packed bytes) per call, 8 even + 8 odd lanes. +CPU_MOE_TARGET("avx2,fma") +inline __m256 nvfp4_blk_avx2(const uint8_t* pk, const float* xeb, const float* xob, + __m256 mag8, float sc) { + __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); + __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); + __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); + __m256 prod = _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); + return _mm256_mul_ps(prod, _mm256_set1_ps(sc)); +} + +CPU_MOE_TARGET("avx2,fma") +float dot_nvfp4_avx2(const uint8_t* packed, const uint8_t* scale, float global, + const float* xe, const float* xo, int K, const float* e2m1, + const float* e4m3) { + const __m256 mag8 = _mm256_loadu_ps(e2m1); // e2m1[0..7] magnitudes + __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); + const int nb = K / 16; + int b = 0; + for (; b + 2 <= nb; b += 2) { + acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8, + xo + (size_t)b * 8, mag8, e4m3[scale[b]])); + acc1 = _mm256_add_ps(acc1, nvfp4_blk_avx2(packed + (size_t)(b + 1) * 8, xe + (size_t)(b + 1) * 8, + xo + (size_t)(b + 1) * 8, mag8, e4m3[scale[b + 1]])); + } + for (; b < nb; ++b) + acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8, + xo + (size_t)b * 8, mag8, e4m3[scale[b]])); + return hsum256(_mm256_add_ps(acc0, acc1)) * global; +} + +// AVX-VNNI W4A8: decode 8 packed bytes (16 nibbles) of one 16-block to int8 [lo(8),hi(8)] +// via PSHUFB against the e2m1*2 LUT (replaces the 2 vpermps fp32 expands -- ~4x less +// port-5 traffic). lo=even-K weights, hi=odd-K, matching the [even(8),odd(8)] act layout. +CPU_MOE_TARGET("avx2,avxvnni,fma") +inline __m128i nvfp4_decode_block_i8(const uint8_t* pk, __m128i lut) { + __m128i b = _mm_loadl_epi64(reinterpret_cast(pk)); // 8 bytes + __m128i lo = _mm_and_si128(b, _mm_set1_epi8(0x0F)); + __m128i hi = _mm_and_si128(_mm_srli_epi16(b, 4), _mm_set1_epi8(0x0F)); + return _mm_shuffle_epi8(lut, _mm_unpacklo_epi64(lo, hi)); // [lo(8),hi(8)] -> int8 +} + +// Two 16-blocks per VPDPBUSD (32 int8). Sign trick (ggml mul_add_epi8): |w|*(sign(w)*a)=w*a, +// so u8*s8 needs no offset/correction term. Per-block scale (e4m3 * act-scale) folded in fp32 +// (lanes 0-3 -> block b, 4-7 -> block b+1). Bit-faithful weight; only the int8 activation quant +// (W4A8) differs from the bf16 reference. +CPU_MOE_TARGET("avx2,avxvnni,fma") +float dot_nvfp4_i8_vnni(const uint8_t* packed, const uint8_t* scale, float global, + const int8_t* asi8, int K, const float* e4m3, const float* asb) { + const __m128i lut = _mm_loadu_si128(reinterpret_cast(kE2M1x2)); + __m256 accF = _mm256_setzero_ps(); + const int nb = K / 16; + int b = 0; + for (; b + 2 <= nb; b += 2) { + __m128i wb = nvfp4_decode_block_i8(packed + (size_t)b * 8, lut); + __m128i wb1 = nvfp4_decode_block_i8(packed + (size_t)(b + 1) * 8, lut); + __m256i w = _mm256_set_m128i(wb1, wb); // [blk b | blk b+1] + __m256i a = _mm256_loadu_si256(reinterpret_cast(asi8 + (size_t)b * 16)); + __m256i aw = _mm256_sign_epi8(w, w); // |w| (u8 operand) + __m256i sa = _mm256_sign_epi8(a, w); // sign(w)*a (s8 operand) + __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); + __m256 scv = _mm256_blend_ps(_mm256_set1_ps(e4m3[scale[b]] * asb[b]), + _mm256_set1_ps(e4m3[scale[b + 1]] * asb[b + 1]), 0xF0); + accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), scv, accF); + } + float s = hsum256(accF); + for (; b < nb; ++b) { // tail (odd block count) + const uint8_t* pk = packed + (size_t)b * 8; + const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; + int isum = 0; + for (int j = 0; j < 8; ++j) + isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; + s += (e4m3[scale[b]] * asb[b]) * (float)isum; + } + return s * (0.5f * global); +} + +#ifdef CPU_MOE_HAS_AVX512VNNI + +// Software-prefetch distance for the W4A8 weight stream, in 16-K blocks (8 packed +// bytes each). Returns -1 when FREETOKEN_CPU_MOE_PF_BLOCKS is unset: the kernel then +// uses the built-in default min(512 blocks = 4 KB, 2 rows) -- 4 KB is the empirical +// optimum on large-row machines (Emerald Rapids sweep), while the 2-row cap keeps a small-row +// model's overshoot bounded (the executor works in 32-row tiles, so a fixed byte +// distance otherwise prefetches another worker's tile: duplicated DRAM traffic that +// regresses at the bandwidth ceiling). An EXPLICIT env value is honored verbatim +// (no clamp; 0 disables): the per-machine optimum can sit past the safe default (+20% +// at 4 KB on a 24-thread Ice Lake with 256B rows), so the escape hatch must reach it. +// Prefetch never faults, so overshooting a row/bank tail is safe. +static int nvfp4_pf_blocks() { + static const int v = [] { + const char* s = getenv("FREETOKEN_CPU_MOE_PF_BLOCKS"); + return (s && s[0]) ? atoi(s) : -1; + }(); + return v; +} + +// AVX-512 VNNI W4A8: FOUR 16-K blocks per VPDPBUSD (64 int8) -- 2x the AVX-VNNI +// (256-bit) path. Decode 32 packed bytes -> 64 int8 with a single _mm512_shuffle_epi8 +// (e2m1*2 LUT replicated to all 4 128-bit lanes). AVX-512 has no _mm512_sign_epi8, so +// the u8*s8 sign trick (|w| as u8, sign(w)*a as s8) uses abs_epi8 + a masked negate. +// Bit-faithful weight; only the int8 activation quant (W4A8) differs from bf16. +// One 4-block group (64 int8) -> scaled fp32 partial (16 lanes). Isolated as a helper so +// the caller can run several independent chains into separate accumulators (the decode -> +// dpbusd -> scale chain is long, so a single accumulator leaves the core latency-bound). +CPU_MOE_TARGET("avx512f,avx512bw,avx512vnni,avx2") +static inline __m512 nvfp4_i8_grp4(const uint8_t* packed, const uint8_t* scale, + const int8_t* asi8, const float* e4m3, const float* asb, + int b, __m512i lut, __m512i idx, __m512i mask0F, + __m512i idxsc) { + const __mmask64 hi_half = 0xFF00FF00FF00FF00ULL; // bytes 8-15 of each 128b lane -> hi nibbles + __m256i raw = _mm256_loadu_si256(reinterpret_cast(packed + (size_t)b * 8)); + __m512i src = _mm512_permutexvar_epi64(idx, _mm512_castsi256_si512(raw)); + __m512i lo = _mm512_and_si512(src, mask0F); + __m512i hi = _mm512_and_si512(_mm512_srli_epi16(src, 4), mask0F); + __m512i comb = _mm512_mask_blend_epi8(hi_half, lo, hi); + __m512i w = _mm512_shuffle_epi8(lut, comb); // 64 int8 weights (e2m1*2) + // u8*s8 sign trick without _mm512_sign_epi8: aw=|w|, sa = (w<0 ? -a : a) (a moot at w==0) + __m512i a = _mm512_loadu_si512(reinterpret_cast(asi8 + (size_t)b * 16)); + __m512i aw = _mm512_abs_epi8(w); + __mmask64 neg = _mm512_movepi8_mask(w); + __m512i sa = _mm512_mask_sub_epi8(a, neg, _mm512_setzero_si512(), a); + __m512i di = _mm512_dpbusd_epi32(_mm512_setzero_si512(), aw, sa); // 16 int32 (groups of 4) + // int32[0..3]->blk b, [4..7]->b+1, [8..11]->b+2, [12..15]->b+3. The 4 block scales + // (e4m3 LUT x per-block act scale) are computed vectorized: a 4-byte load + epu8->epi32 + // widen + one 4-lane LUT gather + one mul replaces 8 scalar loads + 4 scalar muls + + // a set_ps assembly, which otherwise dominates the per-group op count. + int sc_raw; + memcpy(&sc_raw, scale + b, 4); + __m128i sc4 = _mm_cvtepu8_epi32(_mm_cvtsi32_si128(sc_raw)); + __m128 s4 = _mm_mul_ps(_mm_i32gather_ps(e4m3, sc4, 4), _mm_loadu_ps(asb + b)); + __m512 scv = _mm512_permutexvar_ps(idxsc, _mm512_castps128_ps512(s4)); + return _mm512_mul_ps(_mm512_cvtepi32_ps(di), scv); +} + +CPU_MOE_TARGET("avx512f,avx512bw,avx512vnni,avx2") +float dot_nvfp4_i8_avx512vnni(const uint8_t* packed, const uint8_t* scale, float global, + const int8_t* asi8, int K, const float* e4m3, const float* asb) { + const __m512i lut = _mm512_broadcast_i32x4( + _mm_loadu_si128(reinterpret_cast(kE2M1x2))); + const __m512i idx = _mm512_set_epi64(3, 3, 2, 2, 1, 1, 0, 0); // block i -> 128b lane i + const __m512i mask0F = _mm512_set1_epi8(0x0F); + const __m512i idxsc = _mm512_set_epi32(3, 3, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0); + __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); + __m512 acc2 = _mm512_setzero_ps(), acc3 = _mm512_setzero_ps(); + const int nb = K / 16; + // Built-in default: min(4 KB, 2 rows) -- see nvfp4_pf_blocks(). An explicitly set + // env value is used verbatim so a per-machine sweep can reach operating points past + // the conservative default. + const int pfb = nvfp4_pf_blocks(); + const int pf = (pfb < 0) ? std::min(512, 2 * nb) : pfb; + int b = 0; + // Four independent 4-block groups per iter -> four accumulator chains hide the ~10-op + // decode->dpbusd->scale latency (1-2 chains leave the loop latency-bound: the per-core + // rate sat at ~60% of the core's achievable DRAM stream rate). + for (; b + 16 <= nb; b += 16) { + // Prefetch the weight stream ahead: the interleaved decode lowers the L1-miss + // concurrency the HW prefetcher sustains on its own (134 -> 167 GB/s on Emerald Rapids). + if (pf > 0) { + _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8), + _MM_HINT_T0); + _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8 + 64), + _MM_HINT_T0); + } + acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, + mask0F, idxsc)); + acc1 = _mm512_add_ps(acc1, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 4, lut, idx, + mask0F, idxsc)); + acc2 = _mm512_add_ps(acc2, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 8, lut, idx, + mask0F, idxsc)); + acc3 = _mm512_add_ps(acc3, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 12, lut, idx, + mask0F, idxsc)); + } + for (; b + 4 <= nb; b += 4) + acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, + mask0F, idxsc)); + float s = _mm512_reduce_add_ps( + _mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3))); + for (; b < nb; ++b) { // tail (<4 remaining 16-K blocks) + const uint8_t* pk = packed + (size_t)b * 8; + const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; + int isum = 0; + for (int j = 0; j < 8; ++j) + isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; + s += (e4m3[scale[b]] * asb[b]) * (float)isum; + } + return s * (0.5f * global); +} +#endif // avx512vnni available +#endif + +nvdot_fn select_nvdot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (t >= ISA_AVX512) return dot_nvfp4_avx512; + if (t >= ISA_AVX2) return dot_nvfp4_avx2; +#endif + (void)t; + return dot_nvfp4_scalar; +} + +// Best W4A8 (int8-activation) nvfp4 dot, or nullptr if no SIMD VNNI (caller keeps the +// faithful fp32 nvdot path). The scalar i8 dot exists only as a correctness reference. +nvi8dot_fn select_nvi8dot() { +#if CPU_MOE_X86 +#if defined(CPU_MOE_HAS_AVX512VNNI) + if (cpu_has_avx512vnni()) return dot_nvfp4_i8_avx512vnni; +#endif + if (cpu_has_avxvnni()) return dot_nvfp4_i8_vnni; +#endif + return nullptr; +} + +// Quantize the pre-deinterleaved fp32 even/odd activations to per-16-block int8 in the +// [even(8),odd(8)] layout the VNNI dot expects. Done once per token/route (amortized over +// every output row), so a scalar pass is fine relative to the GEMV. +void quant_i8_pg16(const float* xe, const float* xo, int K, int8_t* asi8, float* asb) { + const int nb = K / 16; + for (int b = 0; b < nb; ++b) { + const float* xeb = xe + (size_t)b * 8; + const float* xob = xo + (size_t)b * 8; + float amax = 0.0f; + for (int j = 0; j < 8; ++j) + amax = std::max(amax, std::max(std::fabs(xeb[j]), std::fabs(xob[j]))); + const float s = amax > 0.0f ? amax / 127.0f : 1.0f; + asb[b] = s; + const float inv = 1.0f / s; + int8_t* ae = asi8 + (size_t)b * 16; + for (int j = 0; j < 8; ++j) { + int qe = (int)std::lround(xeb[j] * inv), qo = (int)std::lround(xob[j] * inv); + ae[j] = (int8_t)std::max(-127, std::min(127, qe)); + ae[8 + j] = (int8_t)std::max(-127, std::min(127, qo)); + } + } +} + +class Nvfp4Format final : public MoeFormat { + public: + Nvfp4Format(const MoeParams& p, const FormatArgs& a) + : MoeFormat(p), nvdot_(select_nvdot()), nvi8dot_(select_nvi8dot()) { + gate_up_ = a.bank("gate_up"); + gu_scale_ = a.bank("gate_up_scale"); + gu_global_ = a.bank("gate_up_global"); + down_ = a.bank("down"); + dn_scale_ = a.bank("down_scale"); + dn_global_ = a.bank("down_global"); + const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k; + xe_.resize(tokens * (p.H / 2)); + xo_.resize(tokens * (p.H / 2)); + ge_.resize(routes * (p.I / 2)); + go_.resize(routes * (p.I / 2)); + if (nvi8dot_ != nullptr) { + xi8_.resize(tokens * p.H); + xas_.resize(tokens * (p.H / 16)); + gi8_.resize(routes * p.I); + gas_.resize(routes * (p.I / 16)); + } + isa_ = tier_name(pick_isa()); + if (nvi8dot_ != nullptr) isa_ += cpu_has_avx512vnni() ? "+avx512vnni(nvfp4-w4a8)" : "+vnni(nvfp4-w4a8)"; + for (int i = 0; i < 16; ++i) e2m1_lut_[i] = kE2M1[i]; + for (int i = 0; i < 256; ++i) e4m3_lut_[i] = e4m3_decode((uint8_t)i); + } + + std::string isa_name() const override { return isa_; } + + + void prep_input(int tok, const bf16_t* x) override { + float* xe = xe_.data() + (size_t)tok * (p.H / 2); + float* xo = xo_.data() + (size_t)tok * (p.H / 2); + deinterleave_bf16_f32(x, xe, xo, p.H); + if (nvi8dot_ != nullptr) + quant_i8_pg16(xe, xo, p.H, xi8_.data() + (size_t)tok * p.H, xas_.data() + (size_t)tok * (p.H / 16)); + } + + bool preps_inter() const override { return true; } + + void prep_inter(int64_t r, bf16_t* g) override { + float* ge = ge_.data() + (size_t)r * (p.I / 2); + float* go = go_.data() + (size_t)r * (p.I / 2); + deinterleave_bf16_f32(g, ge, go, p.I); + if (nvi8dot_ != nullptr) + quant_i8_pg16(ge, go, p.I, gi8_.data() + (size_t)r * p.I, gas_.data() + (size_t)r * (p.I / 16)); + } + + void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float w_in, + bf16_t* g) const override { + const int H = p.H; + const uint8_t* packed = static_cast(tbl_at(gate_up_, layer)); + const uint8_t* scale = static_cast(tbl_at(gu_scale_, layer)); + const uint16_t* global = static_cast(tbl_at(gu_global_, layer)); + const float* xe = xe_.data() + (size_t)tok * (H / 2); + const float* xo = xo_.data() + (size_t)tok * (H / 2); + const int8_t* xi8 = nvi8dot_ != nullptr ? xi8_.data() + (size_t)tok * H : nullptr; + const float* xas = nvi8dot_ != nullptr ? xas_.data() + (size_t)tok * (H / 16) : nullptr; + gated_rows(p, i0, i1, w_in, g, [&](int row) { + const size_t r = (size_t)e * (2 * p.I) + row; + if (nvi8dot_ != nullptr) + return nvi8dot_(packed + r * (size_t)(H / 2), scale + r * (size_t)(H / 16), + fp16_to_f32(global[r]), xi8, H, e4m3_lut_, xas); + return nvdot_(packed + r * (size_t)(H / 2), scale + r * (size_t)(H / 16), + fp16_to_f32(global[r]), xe, xo, H, e2m1_lut_, e4m3_lut_); + }); + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const int I = p.I; + const uint8_t* packed = static_cast(tbl_at(down_, layer)); + const uint8_t* scale = static_cast(tbl_at(dn_scale_, layer)); + const uint16_t* global = static_cast(tbl_at(dn_global_, layer)); + for (int h = h0; h < h1; ++h) { + for (int j = 0; j < n; ++j) { + const size_t r = (size_t)routes[j].r; + const size_t row = (size_t)routes[j].e * p.H + h; + out[j * kRowBlock + h - h0] = nvi8dot_ != nullptr + ? nvi8dot_(packed + row * (size_t)(I / 2), scale + row * (size_t)(I / 16), fp16_to_f32(global[row]), + gi8_.data() + r * I, I, e4m3_lut_, gas_.data() + r * (I / 16)) + : nvdot_(packed + row * (size_t)(I / 2), scale + row * (size_t)(I / 16), fp16_to_f32(global[row]), + ge_.data() + r * (I / 2), go_.data() + r * (I / 2), I, e2m1_lut_, e4m3_lut_); + } + } + } + + private: + nvdot_fn nvdot_; + nvi8dot_fn nvi8dot_; // AVX-VNNI W4A8 nvfp4 dot (nullptr -> use fp32 nvdot) + // ModelOpt rows: packed e2m1 [E,2I,H/2], e4m3 scales per 16 along K [E,2I,H/16], fp16 row globals [E,2I]; down likewise with K = I. + const uint64_t* gate_up_ = nullptr; + const uint64_t* gu_scale_ = nullptr; + const uint64_t* gu_global_ = nullptr; + const uint64_t* down_ = nullptr; + const uint64_t* dn_scale_ = nullptr; + const uint64_t* dn_global_ = nullptr; + std::string isa_; + float e2m1_lut_[16]; + float e4m3_lut_[256]; + // Activations pre-deinterleaved to fp32 even/odd for the row-major dot; with VNNI also per-16-block int8 + scale. + AlignedVector xe_, xo_, ge_, go_; + AlignedVector xi8_, gi8_; + AlignedVector xas_, gas_; +}; + +} // namespace + +std::unique_ptr make_nvfp4_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp new file mode 100644 index 0000000000..600a194ebc --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp @@ -0,0 +1,216 @@ +#include +#include + +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// -------------------------------- Q4_0 (W4A8) -------------------------------- +// Native GGUF Q4_0 experts (gemma4 GGUF): per-32 block = fp16 scale d + 16 packed +// bytes; byte j holds element j in its low nibble and j+16 in its high nibble, so a +// block's storage order is [lo0..lo15, hi0..hi15] and w = (nibble - 8) * d. Matches +// the reference dequant (models/gguf/dequant.py) and the packed banks the GPU offload +// path streams. +// +// llama.cpp ggml_vec_dot_q4_0_q8_0: W4A8. The activation is pre-quantized to Q8_0 +// (per-32-block int8 ``aq`` + fp32 scale ``asb``); each block unpacks its 16 bytes to +// 32 int8 weights in [-8,7] (bytes_from_nibbles_32: low nibbles -> elems 0..15, high +// -> 16..31) and runs an integer block dot -- VPDPBUSD (AVX-VNNI) or VPMADDUBSW+VPMADDWD +// (AVX2) with the ggml sign trick |w|*(sign(w)*a)=w*a, or a scalar int loop -- then +// scales the block sum by wd*xd in fp32. No fp weight dequant / shuffle chain. The GPU +// offload path (ggml_moe_a8_vec / MMVQ) is also W4A8, so cpu and hybrid stay close. +using q4dot_fn = float (*)(const uint8_t*, const int8_t*, const float*, int); + +float q4_0_dot_i8_scalar(const uint8_t* w, const int8_t* aq, const float* asb, int K) { + float acc = 0.0f; + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 18; + uint16_t dh; + std::memcpy(&dh, blk, sizeof(dh)); + const uint8_t* q = blk + 2; // 16 nibble bytes + const int8_t* a = aq + (size_t)b * 32; + int isum = 0; + for (int j = 0; j < 16; ++j) { + isum += ((int)(q[j] & 0x0F) - 8) * (int)a[j]; // elem j + isum += ((int)(q[j] >> 4) - 8) * (int)a[16 + j]; // elem 16+j + } + acc += fp16_to_f32(dh) * asb[b] * (float)isum; + } + return acc; +} + +#if CPU_MOE_X86 +// fp16 block scale -> fp32 via HW F16C (single value in lane 0). +CPU_MOE_TARGET("f16c") +static inline float q4_scale(uint16_t h) { + return _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128((int)h))); +} + +// Unpack one Q4_0 block's 16 bytes -> 32 int8 weights in [-8,7] (elems 0..15 = low +// nibbles, 16..31 = high nibbles). ``eight`` = _mm256_set1_epi8(8). +CPU_MOE_TARGET("avx2") +static inline __m256i q4_unpack32(const uint8_t* blk, __m128i mask, __m256i eight) { + const __m128i qb = _mm_loadu_si128(reinterpret_cast(blk + 2)); + const __m128i lo = _mm_and_si128(qb, mask); + const __m128i hi = _mm_and_si128(_mm_srli_epi16(qb, 4), mask); + return _mm256_sub_epi8(_mm256_set_m128i(hi, lo), eight); +} + +// AVX2 W4A8 (llama.cpp non-VNNI mul_sum_i8_pairs): integer block dot via VPMADDUBSW + +// VPMADDWD (sign trick), scaled by wd*xd. |aw*sa| pair sums <= 8*127*2 < 32767 -> no +// int16 saturation. This is the fast path on AVX2 CPUs without AVX-VNNI (and the +// avx512-tier fallback, since the block dot is 256-bit either way). +CPU_MOE_TARGET("avx2,fma,f16c") +float q4_0_dot_i8_avx2(const uint8_t* w, const int8_t* aq, const float* asb, int K) { + const __m128i mask = _mm_set1_epi8(0x0F); + const __m256i eight = _mm256_set1_epi8(8); + const __m256i ones16 = _mm256_set1_epi16(1); + __m256 accF = _mm256_setzero_ps(); + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 18; + _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); + uint16_t dh; + std::memcpy(&dh, blk, sizeof(dh)); + __m256i wq = q4_unpack32(blk, mask, eight); + __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); + __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) + __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) + __m256i d32 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, sa), ones16); // 8 int32 + accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d32), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); + } + return hsum256(accF); +} + +// AVX-VNNI W4A8: one VPDPBUSD per block (the fast path on modern CPUs). +CPU_MOE_TARGET("avx2,avxvnni,fma,f16c") +float q4_0_dot_i8_vnni(const uint8_t* w, const int8_t* aq, const float* asb, int K) { + const __m128i mask = _mm_set1_epi8(0x0F); + const __m256i eight = _mm256_set1_epi8(8); + __m256 accF = _mm256_setzero_ps(); + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 18; + _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); + uint16_t dh; + std::memcpy(&dh, blk, sizeof(dh)); + __m256i wq = q4_unpack32(blk, mask, eight); + __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); + __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) + __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) + __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); + // All 32 elems of the block share wd*xd; distribute over di's 8 partial sums and + // reduce at the end (equivalent to scale * block_total). + accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); + } + return hsum256(accF); +} +#endif // CPU_MOE_X86 + +// All tiers are W4A8 (int8 activations pre-quantized to Q8_0). AVX-VNNI is orthogonal to +// the ISA tier (gated by cpu_has_avxvnni() / FREETOKEN_CPU_MOE_NO_VNNI), so it wins when +// present; otherwise the 256-bit VPMADDUBSW kernel covers both the avx2 and avx512 tiers. +q4dot_fn select_q4dot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (cpu_has_avxvnni()) return q4_0_dot_i8_vnni; + if (t >= ISA_AVX2) return q4_0_dot_i8_avx2; +#endif + (void)t; + return q4_0_dot_i8_scalar; +} + +// Quantize a bf16 activation row to Q8_0 (llama.cpp): per-32-block symmetric int8 in +// natural order + one fp32 scale (amax/127) per block. Done once per token/route, +// amortized over every output row that the q4_0 W4A8 GEMV reads. K % 32 == 0. +void quant_q8_0(const bf16_t* x, int K, int8_t* aq, float* asb) { + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const bf16_t* xb = x + (size_t)b * 32; + float xf[32], amax = 0.0f; + for (int j = 0; j < 32; ++j) { + xf[j] = bf16_to_f32(xb[j]); + amax = std::max(amax, std::fabs(xf[j])); + } + const float d = amax > 0.0f ? amax / 127.0f : 1.0f; + asb[b] = d; + const float inv = amax > 0.0f ? 1.0f / d : 0.0f; + int8_t* o = aq + (size_t)b * 32; + for (int j = 0; j < 32; ++j) + o[j] = (int8_t)std::max(-127, std::min(127, (int)std::lround(xf[j] * inv))); + } +} + +class Q4_0Format final : public MoeFormat { + public: + Q4_0Format(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), dot_(select_q4dot()) { + gate_up_ = a.bank("gate_up"); + down_ = a.bank("down"); + const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k; + xi8_.resize(tokens * p.H); + xas_.resize(tokens * (p.H / 32)); + gi8_.resize(routes * p.I); + gas_.resize(routes * (p.I / 32)); + if (p.H % 32 != 0 || p.I % 32 != 0) + throw std::runtime_error("Q4_0 CPU MoE requires H and I to be multiples of 32"); + gu_row_bytes_ = (p.H / 32) * 18; // K = H (gate_up rows) + dn_row_bytes_ = (p.I / 32) * 18; // K = I (down rows) + // select_q4dot picks VPDPBUSD / VPMADDUBSW / scalar for the tier, so the tag reflects which of those it resolved to. + isa_ = std::string(tier_name(pick_isa())) + (cpu_has_avxvnni() ? "+vnni(q4_0-w4a8)" : "+q4_0-w4a8"); + } + + std::string isa_name() const override { return isa_; } + + + void prep_input(int tok, const bf16_t* x) override { + quant_q8_0(x, p.H, xi8_.data() + (size_t)tok * p.H, xas_.data() + (size_t)tok * (p.H / 32)); + } + + bool preps_inter() const override { return true; } + + void prep_inter(int64_t r, bf16_t* g) override { + quant_q8_0(g, p.I, gi8_.data() + (size_t)r * p.I, gas_.data() + (size_t)r * (p.I / 32)); + } + + void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float w_in, + bf16_t* g) const override { + const uint8_t* packed = static_cast(tbl_at(gate_up_, layer)); + const int8_t* xi8 = xi8_.data() + (size_t)tok * p.H; + const float* xas = xas_.data() + (size_t)tok * (p.H / 32); + gated_rows(p, i0, i1, w_in, g, [&](int row) { + return dot_(packed + ((size_t)e * (2 * p.I) + row) * (size_t)gu_row_bytes_, xi8, xas, p.H); + }); + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const uint8_t* packed = static_cast(tbl_at(down_, layer)); + for (int h = h0; h < h1; ++h) { + for (int j = 0; j < n; ++j) { + const size_t r = (size_t)routes[j].r; + out[j * kRowBlock + h - h0] = dot_(packed + ((size_t)routes[j].e * p.H + h) * (size_t)dn_row_bytes_, + gi8_.data() + r * p.I, gas_.data() + r * (p.I / 32), p.I); + } + } + } + + private: + q4dot_fn dot_; + // GGUF Q4_0 rows of 32-element blocks (fp16 scale + 16 nibble bytes): gate_up [E,2I,H/32*18], down [E,H,I/32*18]. + const uint64_t* gate_up_ = nullptr; + const uint64_t* down_ = nullptr; + int gu_row_bytes_ = 0, dn_row_bytes_ = 0; + std::string isa_; + // Q8_0 activations: per-32-block int8 + fp32 scale, for the input and the intermediate. + AlignedVector xi8_, gi8_; + AlignedVector xas_, gas_; +}; + +} // namespace + +std::unique_ptr make_q4_0_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py index e70cad31d3..9430ab2a84 100644 --- a/python/freetoken/moe/benchbw.py +++ b/python/freetoken/moe/benchbw.py @@ -490,17 +490,16 @@ def _cpu_moe_step_fns(ex, wl: Workload, E: int): disjoint experts (working set >> LLC) -- the DRAM-bound regime hybrid decode actually runs in; ``run()`` executes one step (releases the GIL inside the extension).""" bs, top_k = 1, wl.top_k - io = ex._io_for(bs) + io = ex._io io["x"].copy_(torch.randn(bs, wl.hidden, dtype=torch.bfloat16) * 0.1) io["w"].copy_(torch.rand(bs, top_k, dtype=torch.float32)) - task = ex._task_for(0, bs) base = torch.arange(top_k, dtype=torch.int32) def set_ids(i: int) -> None: io["ids"].copy_((((i * top_k) + base) % E).view(bs, top_k)) def run() -> None: - ex._ext.run_task(task) + ex._ext.run(0, bs) return set_ids, run @@ -566,7 +565,7 @@ def measure_overlap_bw(fmt: str, wl: Workload, device: torch.device, backend its bandwidth-matched fetch split: fetched : cpu-computed = pcie_ov : cpu_ov. Both sides hammer their own pinned banks flat out for ``seconds`` from a shared - barrier: a worker thread loops bs=1 CPU decode steps (``run_task`` releases the GIL) + barrier: a worker thread loops bs=1 CPU decode steps (``run`` releases the GIL) while the main thread loops full-layer ``copy_missing`` gathers (synchronized per copy so the DMA is really in flight, not just enqueued). Each side reports bytes / its own elapsed; the two windows differ by at most one CPU step + one gather. @@ -663,7 +662,7 @@ def _bench_format(fmt: str, wl: Workload, device: torch.device, threshold: float entry["isa_sweep"] = c["isa_sweep"] # The tier label comes from the bf16 dot chooser, so it can over-report tiers # the actual per-format kernel doesn't distinguish. Disclose the real story. - if fmt == "nvfp4": # WF_NVFP4 is the only use_vnni format + if fmt == "nvfp4": # nvfp4 is the only use_vnni format _note(entry, "isa tier labels are nominal: this W4A8 kernel rides AVX-VNNI " "regardless of the forced tier") elif fmt in ("mxfp4_triton", "ds_fp4", "fp8_block"): diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index f519687ed6..079c1f43d3 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -1,18 +1,13 @@ """Python wrapper around the ``_cpu_moe`` C++ executor (--moe-strategy cpu). -Owns the persistent CPU worker pool, the per-batch-size pinned IO buffers, and -the per-(layer, batch-size) host-func task descriptors. ``decode`` issues the -whole CUDA-graph-capturable sequence on the current stream: +Owns the persistent CPU worker pool and one pinned IO block of ``max_tokens`` rows that +every layer and batch size shares (one task is in flight at a time). ``decode`` issues +the whole CUDA-graph-capturable sequence on the current stream: - D2H (hidden, topk_ids, topk_weights -> pinned) - -> submit host node (cudaLaunchHostFunc: enqueue MoE task to the pool) - -> sync host node (cudaLaunchHostFunc: spin until the pool drains) + D2H (hidden, topk_ids, topk_weights -> pinned IO rows [:bs]) + -> submit (doorbell naming (layer, bs), or a cudaLaunchHostFunc host node) + -> sync (wait for the CPU pool to finish) -> H2D (pinned expert output -> GPU) - -Buffers and tasks are allocated lazily per batch size. GraphRunner runs an eager -``model.forward()`` at each batch size immediately before capturing it, so the -first (eager) call materializes the stable pinned buffers + task pointers that -the subsequent capture embeds in its host/memcpy nodes. """ from __future__ import annotations @@ -35,28 +30,27 @@ # on a 75-layer model). Instead the GPU raises a mapped-pinned "ready" flag at submit; a # persistent CPU coordinator (in _cpu_moe) polls it, runs the layer, and sets a "done" # flag the GPU waits on at sync -- no host-func round-trip. Both GPU-side operations are -# STREAM MEMORY OPERATIONS (cuStreamWriteValue64 / cuStreamWaitValue64, resolved from the -# driver at runtime): they execute on the GPU front end with no SM-resident kernel, so +# STREAM MEMORY OPERATIONS (cuStreamWriteValue64 / cuStreamWaitValue64 on CUDA, see +# gpu_sync.h): they execute on the GPU front end with no SM-resident kernel, so # GPU utilization stays truthful during the CPU compute window. (The first cut used a # spin-wait kernel; that pinned reported utilization at 99% and laptop CPU/GPU dynamic # power schedulers responded by clamping the CPU frequency -- a net decode regression on -# power-coupled edge devices.) Each (layer, decode batch size) pair gets its own flag -# slot, so every captured decode graph rides the handshake. Where memops are unavailable -# (Windows WDDM, vGPU, old drivers -- functionally probed at startup) or the slot -# capacity is exceeded, decode keeps the host-func path (functional, slower). A Python -# watchdog thread turns a wedged coordinator into a loud RuntimeError (via err[] + -# raise_if_unhealthy) instead of an indefinite stream stall. +# power-coupled edge devices.) The doorbell word names the (layer, decode batch size) +# task, so one ready/done pair serves every layer and every captured graph (gpu_sync.h). +# Where memops are unavailable (some driver models, vGPU, old drivers -- functionally probed at +# startup), decode keeps the host-func path (functional, slower). ROCm 7.14 cannot replay +# that path from a graph, so capture fails loudly there instead. A Python watchdog thread +# turns a wedged coordinator into a loud RuntimeError (via raise_if_unhealthy) instead of +# an indefinite stream stall. # Caveat: the coordinator busy-polls one core while decode traffic flows (idle backoff # otherwise); FREETOKEN_CPU_MOE_FLAG_SYNC=0 opts out entirely. _FLAG_SYNC = os.getenv("FREETOKEN_CPU_MOE_FLAG_SYNC", "1") != "0" -# Flag slots per MoE layer: covers this many distinct decode batch sizes (captured graph -# sizes plus any eager padded sizes); more than that is unheard of, and the overflow -# just keeps the host-func path for the extra combos. -_FLAG_SLOTS_PER_LAYER = 16 # Pool threads spin up to 50 ms before they park. FREETOKEN_CPU_MOE_SPIN=0 opts out. _SPIN_WAIT = os.getenv("FREETOKEN_CPU_MOE_SPIN", "1") != "0" +# ds_fp4 runs its input FP8 round-trip on the GPU when it can; tests turn this off to compare against the CPU round-trip. +_DSFP4_GPU_PREQUANT = True -# Activation ids must match ActKind in csrc/cpu_moe/cpu_moe_ext.cpp. Id 3 is the +# Activation ids must match ActKind in csrc/cpu_moe/common.h. Id 3 is the # clamped (up + 1) swiglu: "swigluoai" runs it in the generic GEMV epilogue, # "gpt_oss_swiglu" is the same math fused inside the mxfp4 kernel. _ACT_IDS = { @@ -70,10 +64,6 @@ "swiglu_clamp": 4, } -# Weight-format ids must match WFmt in csrc/cpu_moe/cpu_moe_ext.cpp. -_WFMT_IDS = {"bf16": 0, "nvfp4": 1, "mxfp4_triton": 2, "ds_fp4": 3, "q4_0": 4, - "fp8_block": 5} - def compiled_extension_supports(activation: str) -> bool: """Whether the compiled ``_cpu_moe`` extension can serve ``activation`` @@ -157,7 +147,7 @@ def resolve_threads_and_affinity(requested: int) -> tuple[int, list[int]]: class CpuMoeExecutor: """Decode-time CPU expert compute over an ``OffloadMoeCache``'s host banks - (bf16, nvfp4, mxfp4_triton, ds_fp4, q4_0 or fp8_block — see ``_WFMT_IDS`` / ``_resolve_banks``).""" + (the formats ``_cpu_moe.weight_formats()`` lists — see ``_resolve_banks``).""" def __init__( self, @@ -178,26 +168,20 @@ def __init__( from freetoken.moe.legacy_format import canonical_role fmt = fmt or cache.quant_format - if fmt not in _WFMT_IDS: + if fmt not in _cpu_moe.weight_formats(): raise NotImplementedError( f"--moe-strategy cpu/hybrid computes experts on the CPU and supports " - f"{sorted(_WFMT_IDS)} formats, but this checkpoint's experts are " + f"{sorted(_cpu_moe.weight_formats())} formats, but this checkpoint's experts are " f"{fmt!r}; use --moe-strategy offload (GPU-side dequant) instead." ) if activation not in _ACT_IDS: raise NotImplementedError(f"CPU MoE backend: unsupported activation {activation!r}") + if apply_router_weight_on_input and fmt in ("mxfp4_triton", "ds_fp4"): + raise NotImplementedError(f"CPU MoE backend: {fmt} experts apply the router weight on the output only") # ABI probe: a stale prebuilt _cpu_moe.so accepts newer act ids without # error and silently computes the wrong activation in the generic # epilogue -- fail loudly with the rebuild instruction instead. (mxfp4 # handles its act inside the kernel and predates the marker.) - # A prebuilt .so that predates this format would accept the id and then take - # the wrong dequant branch -- silently wrong output rather than a crash. - if _WFMT_IDS[fmt] > getattr(_cpu_moe, "max_weight_format_id", lambda: 4)(): - raise RuntimeError( - f"the compiled _cpu_moe extension predates the {fmt!r} expert format; " - "rebuild it with `python setup.py build_ext --inplace` (or reinstall " - "the wheel) before serving this checkpoint on the cpu/hybrid backend." - ) if _ACT_IDS[activation] >= 3 and fmt != "mxfp4_triton": supported = getattr(_cpu_moe, "max_generic_act_id", lambda: 2)() if _ACT_IDS[activation] > supported: @@ -218,28 +202,42 @@ def __init__( # The per-layer tensors and their pointer tables must outlive the executor # (C++ holds raw addresses into both). self._banks: list[torch.Tensor] = [] - ptrs, (self.H, self.I) = self._resolve_banks( + banks, options, (self.H, self.I) = self._resolve_banks( {canonical_role(name): per_layer for name, per_layer in cache.bank_sources.items()}, fmt ) + # ds_fp4: the reference FP8 activation round-trip is a scalar per-element chain + # that the C++ side runs single-threaded on the CUDA host-callback thread -- + # straight on the decode critical path (~0.3ms/layer at H=4096, every worker and + # the GPU waiting on it). When a GPU is present we run the numerically identical + # round-trip as a captured GPU kernel BEFORE the D2H (see decode_submit) and tell + # the C++ side to skip its own. Measured on DeepSeek-V4-Flash bs=1 decode: + # 12.85 -> 15.65 tok/s, output bit-identical (tests/moe/test_dsfp4_prequant.py). + self._gpu_prequant = fmt == "ds_fp4" and device.type == "cuda" and _DSFP4_GPU_PREQUANT + # the W4A8 activation round-trip block follows the checkpoint's fp8 block (128 on V4, 32 on V4.1) + self._act_block = act_block + if fmt == "ds_fp4": + assert act_block is not None, "ds_fp4 experts need act_block, the checkpoint's fp8 activation quant block" + options.update(act_block=int(act_block), input_prequant=int(self._gpu_prequant)) + if self._gpu_prequant: + logger.info_rank0( + "ds_fp4: input FP8 round-trip moved to the GPU " + "(bit-identical grid; the CPU-side scalar round-trip is skipped)" + ) + # Decide the flag handshake up front (env + device + a functional stream-memop # probe): its coordinator needs a core of its own, which the auto thread sizing # below reserves (a coordinator time-slicing against the GEMV workers measurably # destabilizes throughput on fully-subscribed boxes). self._flag_sync = _FLAG_SYNC and device.type == "cuda" - self._cpu_moe = _cpu_moe # module ref for the decode-path memop calls - if self._flag_sync: - probe_scratch = alloc_pinned_tensor(1, dtype=torch.int64) - probe_scratch.zero_() - if not _cpu_moe.memops_probe( - torch.cuda.current_stream().cuda_stream, probe_scratch.data_ptr() - ): - logger.info_rank0( - "cpu-moe flag handshake unavailable: CUDA stream memory operations " - "are not supported here (Windows WDDM / vGPU / old driver); using " - "the cudaLaunchHostFunc sync" - ) - self._flag_sync = False + if self._flag_sync and not _cpu_moe.memops_probe(torch.cuda.current_stream().cuda_stream): + logger.info_rank0( + "cpu-moe flag handshake unavailable: stream memory operations failed the " + "startup probe; using the cudaLaunchHostFunc sync" + ) + self._flag_sync = False + # ROCm 7.14 returns wrong results when it replays host-function nodes, so there only the doorbell is graph-safe. + self._graph_capture_safe = self._flag_sync or _cpu_moe.host_func_capture_safe() nthreads, core_ids = resolve_threads_and_affinity(num_threads) coord_core = -1 @@ -250,6 +248,17 @@ def __init__( nthreads -= 1 core_ids = core_ids[:-1] self._coord_core = coord_core + alloc_io = ( + (lambda *shape, dtype: torch.empty(*shape, dtype=dtype)) + if device.type == "cpu" + else alloc_pinned_tensor + ) + self._io = { + "x": alloc_io(self.max_tokens, self.H, dtype=torch.bfloat16), + "ids": alloc_io(self.max_tokens, self.top_k, dtype=torch.int32), + "w": alloc_io(self.max_tokens, self.top_k, dtype=torch.float32), + "y": alloc_io(self.max_tokens, self.H, dtype=torch.bfloat16), + } self._ext = _cpu_moe.CpuMoeExecutor( num_threads=nthreads, num_layers=self.num_layers, @@ -260,11 +269,16 @@ def __init__( max_tokens=self.max_tokens, activation_id=_ACT_IDS[activation], apply_router_weight_on_input=1 if apply_router_weight_on_input else 0, - weight_format=_WFMT_IDS[fmt], + weight_format=fmt, + banks=banks, + options=options, swiglu_alpha=float(swiglu_alpha), swiglu_limit=float(swiglu_limit) if swiglu_limit is not None else float("inf"), core_ids=core_ids, - **ptrs, + x_ptr=self._io["x"].data_ptr(), + ids_ptr=self._io["ids"].data_ptr(), + w_ptr=self._io["w"].data_ptr(), + y_ptr=self._io["y"].data_ptr(), ) self.num_threads = nthreads self.core_ids = core_ids @@ -281,34 +295,15 @@ def __init__( ) torch.set_num_threads(clamp) - self._io: dict[int, dict[str, torch.Tensor]] = {} - self._tasks: dict[tuple[int, int], int] = {} - - # Flag-based handshake: mapped-pinned ready/done/err int64 arrays (one slot per - # (MoE layer, decode batch size) pair, allocated as tasks are created) + a - # persistent CPU coordinator that polls ready[], runs the slot's task on the - # pool, and sets done[]. Binary per-step protocol (GPU memops: done=0, ready=1; - # coordinator: consume ready, run, done=1; GPU waits done>=1 -- the WAIT - # immediate is constant, so CUDA-graph replays are safe). err[] is raised by the - # WATCHDOG thread when a ready flag stays unanswered (dead coordinator): it - # poisons done to unblock the stream and raise_if_unhealthy() turns the step - # into a loud error instead of silent stale output. Buffers are kept alive on - # self so the coordinator's pinned pointers stay valid for the executor's - # lifetime (flag_sync itself was decided above, before thread sizing). - self._ready = self._done = self._err = None - self._flag_slots: dict[tuple[int, int], int] = {} # (layer_id, bs) -> slot - self._flag_capacity = self.num_layers * _FLAG_SLOTS_PER_LAYER + # Flag-based handshake: the extension owns the ready/done words (gpu_sync.h) and a + # persistent CPU coordinator that polls the doorbell, runs the task it names on the + # pool, and sets done for the GPU. The WATCHDOG thread marks the executor unhealthy + # when a doorbell stays unanswered (dead coordinator): it poisons done to unblock the + # stream and raise_if_unhealthy() turns the step into a loud error instead of silent + # stale output (flag_sync itself was decided above, before thread sizing). + self._unhealthy = False if self._flag_sync: - self._ready = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) - self._done = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) - self._err = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) - self._ready.zero_() - self._done.zero_() - self._err.zero_() - self._ext.start_flag_coordinator( - self._ready.data_ptr(), self._done.data_ptr(), self._flag_capacity, - self._coord_core, - ) + self._ext.start_flag_coordinator(self._coord_core) self._watchdog_stop = False # The thread target holds a WEAKREF and re-derefs it each tick: a bound # method would strong-reference the executor forever (the loop never ends @@ -325,25 +320,6 @@ def __init__( ) self._watchdog.start() - # ds_fp4: the reference FP8 activation round-trip is a scalar per-element chain - # that the C++ side runs single-threaded on the CUDA host-callback thread -- - # straight on the decode critical path (~0.3ms/layer at H=4096, every worker and - # the GPU waiting on it). When a GPU is present we run the numerically identical - # round-trip as a captured GPU kernel BEFORE the D2H (see decode_submit) and tell - # the C++ side to skip its own. Measured on DeepSeek-V4-Flash bs=1 decode: - # 12.85 -> 15.65 tok/s, output bit-identical (tests/moe/test_dsfp4_prequant.py). - self._gpu_prequant = fmt == "ds_fp4" and device.type == "cuda" - # the W4A8 activation round-trip block follows the checkpoint's fp8 block (128 on V4, 32 on V4.1) - self._act_block = act_block - if fmt == "ds_fp4": - assert act_block is not None, "ds_fp4 experts need act_block, the checkpoint's fp8 activation quant block" - self._ext.set_act_block(act_block) - if self._gpu_prequant: - self._ext.set_input_prequant(True) - logger.info_rank0( - "ds_fp4: input FP8 round-trip moved to the GPU " - "(bit-identical grid; the CPU-side scalar round-trip is skipped)" - ) logger.info_rank0( f"CPU MoE executor ready: threads={nthreads} (pinned to cores " @@ -368,7 +344,11 @@ def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor: self._banks.extend(layers) return table - def _resolve_fp8_block_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: + def _tables(self, **roles: list[torch.Tensor]) -> dict[str, int]: + """Per-layer table addresses keyed by canonical bank role.""" + return {role: self._make_table(layers).data_ptr() for role, layers in roles.items()} + + def _resolve_fp8_block_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]: """Block-FP8 schema: row-major e4m3 + one bf16 scale per 128x128 block; the scale banks pad their last dim, so C++ gets the real row stride.""" gu, gus = banks["gate_up"], banks["gate_up_scale"] dn, dns = banks["down"], banks["down_scale"] @@ -380,28 +360,19 @@ def _resolve_fp8_block_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: assert tuple(dn[0].shape[1:]) == (H, I), (dn[0].shape, H, I) assert gus[0].shape[1] == (2 * I + 127) // 128 and gus[0].shape[2] >= (H + 127) // 128, (gus[0].shape, H, I) assert dns[0].shape[1] == (H + 127) // 128 and dns[0].shape[2] >= (I + 127) // 128, (dns[0].shape, H, I) - ptrs = dict( - gate_up_ptr=self._make_table(gu).data_ptr(), - down_ptr=self._make_table(dn).data_ptr(), - gate_up_scale_ptr=self._make_table(gus).data_ptr(), - gate_up_global_ptr=0, - down_scale_ptr=self._make_table(dns).data_ptr(), - down_global_ptr=0, - gate_up_bias_ptr=0, - down_bias_ptr=0, - fp8_block_gu_scale_stride=int(gus[0].stride(1)), - fp8_block_dn_scale_stride=int(dns[0].stride(1)), - ) - return ptrs, (H, I) + banks = self._tables(gate_up=gu, down=dn, gate_up_scale=gus, down_scale=dns) + options = dict(gu_scale_stride=int(gus[0].stride(1)), dn_scale_stride=int(dns[0].stride(1))) + return banks, options, (H, I) - def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]: - """Return (pointer kwargs for the C++ ctor, (H, I)) for the given format. + def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, dict, tuple[int, int]]: + """Return (banks, options, (H, I)) for the C++ ctor: banks maps each canonical role + the format reads to its per-layer table's address, options holds its integer settings. ``banks[name]`` is a list of ``num_layers`` ``[num_experts, ...]`` tensors (the per-layer host bank contract); shapes are read from the first layer so - per-partition (TP) sizes are exact. Unused pointers are 0. Every pointer kwarg - is actually a per-layer table's address (see ``_make_table``), not a single - bank's -- the C++ ctor resolves ``tbl[layer_id]`` per task. + per-partition (TP) sizes are exact. Every bank address is a per-layer table + (see ``_make_table``), not a single bank's -- the C++ side resolves + ``tbl[layer_id]`` per task. """ if fmt == "bf16": gate_up = banks["gate_up"] @@ -414,17 +385,8 @@ def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]: I = int(gate_up[0].shape[1] // 2) assert gate_up[0].shape[1] == 2 * I assert tuple(down[0].shape[1:]) == (H, I), (down[0].shape, H, I) - ptrs = dict( - gate_up_ptr=self._make_table(gate_up).data_ptr(), - down_ptr=self._make_table(down).data_ptr(), - gate_up_scale_ptr=0, - gate_up_global_ptr=0, - down_scale_ptr=0, - down_global_ptr=0, - gate_up_bias_ptr=0, - down_bias_ptr=0, - ) - return ptrs, (H, I) + banks = self._tables(gate_up=gate_up, down=down) + return banks, {}, (H, I) if fmt == "q4_0": return self._resolve_q4_0_banks(banks) @@ -452,19 +414,10 @@ def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]: assert tuple(gus[0].shape[1:]) == (2 * I, H // 16), (gus[0].shape, I, H) assert tuple(dns[0].shape[1:]) == (H, I // 16), (dns[0].shape, H, I) assert tuple(gug[0].shape[1:]) == (2 * I,) and tuple(dng[0].shape[1:]) == (H,) - ptrs = dict( - gate_up_ptr=self._make_table(gup).data_ptr(), - down_ptr=self._make_table(dnp).data_ptr(), - gate_up_scale_ptr=self._make_table(gus).data_ptr(), - gate_up_global_ptr=self._make_table(gug).data_ptr(), - down_scale_ptr=self._make_table(dns).data_ptr(), - down_global_ptr=self._make_table(dng).data_ptr(), - gate_up_bias_ptr=0, - down_bias_ptr=0, - ) - return ptrs, (H, I) + banks = self._tables(gate_up=gup, down=dnp, gate_up_scale=gus, gate_up_global=gug, down_scale=dns, down_global=dng) + return banks, {}, (H, I) - def _resolve_q4_0_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: + def _resolve_q4_0_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]: """Native GGUF Q4_0 schema (gemma4 GGUF): per-32 blocks (fp16 scale + 16 nibble bytes), row-major over K -- the *same* packed banks the GPU offload path streams. gate_up is [S, 2I, H//32*18], down is [S, H, I//32*18]; the C++ W4A16 GEMV reads a @@ -479,19 +432,10 @@ def _resolve_q4_0_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: assert H % 32 == 0 and I % 32 == 0, (H, I) assert int(gate_up[0].shape[2]) == (H // 32) * 18, (gate_up[0].shape, H) assert int(down[0].shape[2]) == (I // 32) * 18, (down[0].shape, I) - ptrs = dict( - gate_up_ptr=self._make_table(gate_up).data_ptr(), - down_ptr=self._make_table(down).data_ptr(), - gate_up_scale_ptr=0, - gate_up_global_ptr=0, - down_scale_ptr=0, - down_global_ptr=0, - gate_up_bias_ptr=0, - down_bias_ptr=0, - ) - return ptrs, (H, I) + banks = self._tables(gate_up=gate_up, down=down) + return banks, {}, (H, I) - def _resolve_mxfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: + def _resolve_mxfp4_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]: """gpt-oss mxfp4 ``mxfp4_triton`` schema: transposed split-K blocks/scales (N innermost) + per-output-row biases. The C++ kernel streams K and accumulates a contiguous N-block, so the GPU-tiled layout is read in place @@ -510,19 +454,10 @@ def _resolve_mxfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: assert tuple(gus[0].shape[1:]) == (H // 32, 2 * I), (gus[0].shape, H, I) assert tuple(dns[0].shape[1:]) == (I // 32, H), (dns[0].shape, H, I) assert tuple(gob[0].shape[1:]) == (2 * I,) and tuple(dob[0].shape[1:]) == (H,) - ptrs = dict( - gate_up_ptr=self._make_table(gub).data_ptr(), - down_ptr=self._make_table(dnb).data_ptr(), - gate_up_scale_ptr=self._make_table(gus).data_ptr(), - gate_up_global_ptr=0, - down_scale_ptr=self._make_table(dns).data_ptr(), - down_global_ptr=0, - gate_up_bias_ptr=self._make_table(gob).data_ptr(), - down_bias_ptr=self._make_table(dob).data_ptr(), - ) - return ptrs, (H, I) + banks = self._tables(gate_up=gub, down=dnb, gate_up_scale=gus, down_scale=dns, gate_up_bias=gob, down_bias=dob) + return banks, {}, (H, I) - def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: + def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]: """DeepSeek-V4 ``ds_fp4`` schema: row-major e2m1 (2/byte) + e8m0 per-32 block scales, no global, no bias. Layout matches nvfp4 (K contiguous per output row), so the C++ GEMV reads it in place. The kernel additionally FP8-round-trips the @@ -538,57 +473,8 @@ def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: assert tuple(dnp[0].shape[1:]) == (H, I // 2), (dnp[0].shape, H, I) assert tuple(gus[0].shape[1:]) == (2 * I, H // 32), (gus[0].shape, I, H) assert tuple(dns[0].shape[1:]) == (H, I // 32), (dns[0].shape, H, I) - ptrs = dict( - gate_up_ptr=self._make_table(gup).data_ptr(), - down_ptr=self._make_table(dnp).data_ptr(), - gate_up_scale_ptr=self._make_table(gus).data_ptr(), - gate_up_global_ptr=0, - down_scale_ptr=self._make_table(dns).data_ptr(), - down_global_ptr=0, - gate_up_bias_ptr=0, - down_bias_ptr=0, - ) - return ptrs, (H, I) - - def _io_for(self, bs: int) -> dict[str, torch.Tensor]: - io = self._io.get(bs) - if io is None: - alloc_io = ( - (lambda *shape, dtype: torch.empty(*shape, dtype=dtype)) - if self.device.type == "cpu" - else alloc_pinned_tensor - ) - io = { - "x": alloc_io(bs, self.H, dtype=torch.bfloat16), - "ids": alloc_io(bs, self.top_k, dtype=torch.int32), - "w": alloc_io(bs, self.top_k, dtype=torch.float32), - "y": alloc_io(bs, self.H, dtype=torch.bfloat16), - } - self._io[bs] = io - return io - - def _task_for(self, layer_id: int, bs: int) -> int: - key = (layer_id, bs) - task = self._tasks.get(key) - if task is None: - io = self._io_for(bs) - task = self._ext.create_task( - layer_id, - bs, - io["x"].data_ptr(), - io["ids"].data_ptr(), - io["w"].data_ptr(), - io["y"].data_ptr(), - ) - self._tasks[key] = task - # Allocate this (layer, bs) combo a flag slot and register its task with the - # coordinator. Combos past the slot capacity keep the host-func path. - if self._flag_sync and key not in self._flag_slots: - slot = len(self._flag_slots) - if slot < self._flag_capacity: - self._flag_slots[key] = slot - self._ext.register_flag_task(slot, task) - return task + banks = self._tables(gate_up=gup, down=dnp, gate_up_scale=gus, down_scale=dns) + return banks, {}, (H, I) def decode( self, @@ -621,7 +507,14 @@ def decode_submit( :meth:`decode_sync`. The output tensor is allocated here so it stays live (and distinct from the interleaved GPU work) across the overlap window.""" bs = hidden_states.shape[0] - io = self._io_for(bs) + if bs > self.max_tokens: + raise ValueError(f"CPU MoE decode batch {bs} exceeds the executor's max_tokens {self.max_tokens}") + if not self._graph_capture_safe and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "CPU MoE decode cannot be captured into a CUDA graph here: the stream-memop " + "doorbell is off or failed its startup probe, and the host-function sync is not " + "replay-safe on this runtime; run with --cuda-graph-max-bs 0" + ) if self._gpu_prequant: # DSV4: apply the reference FP8 round-trip on the GPU (the same kernel the @@ -632,90 +525,70 @@ def decode_submit( hidden_states = act_quant_fp8_roundtrip(hidden_states, block=self._act_block) # D2H: ship this step's activations + routing to pinned host memory. - io["x"].copy_(hidden_states, non_blocking=True) - io["ids"].copy_(topk_ids.to(torch.int32), non_blocking=True) - io["w"].copy_(topk_weights.to(torch.float32), non_blocking=True) + io = self._io + io["x"][:bs].copy_(hidden_states, non_blocking=True) + io["ids"][:bs].copy_(topk_ids.to(torch.int32), non_blocking=True) + io["w"][:bs].copy_(topk_weights.to(torch.float32), non_blocking=True) - task = self._task_for(layer_id, bs) out = torch.empty_like(hidden_states) - slot = self._flag_slots.get((layer_id, bs)) if self._flag_sync else None - if slot is not None: - # Front-end memops: done[slot]=0 then ready[slot]=1 (the coordinator's - # doorbell). No kernel launched; no host-func round trip. - self._cpu_moe.memop_submit( - torch.cuda.current_stream().cuda_stream, - self._done.data_ptr(), self._ready.data_ptr(), slot, - ) + stream = torch.cuda.current_stream().cuda_stream + if self._flag_sync: + # Front-end memops: reset done, then ring the doorbell with (layer, bs). + # No kernel launched; no host-func round trip. + self._ext.flag_signal(stream, layer_id, bs) else: - stream = torch.cuda.current_stream().cuda_stream - self._ext.submit_with_cuda_stream(stream, task) - return (bs, task, out, slot) + self._ext.submit_with_cuda_stream(stream, layer_id, bs) + return (bs, out) def decode_sync(self, pending: tuple) -> torch.Tensor: """Issue the CPU-pool sync + the H2D result copy for a prior :meth:`decode_submit`, and return the GPU output tensor. With flag-sync the wait is a front-end stream - memop on done[slot] (set by the CPU coordinator); otherwise a cudaLaunchHostFunc.""" - bs, task, out, slot = pending - if slot is not None: - # Front-end WAIT(done[slot] >= 1): blocks this stream's later nodes without + memop on done (set by the CPU coordinator); otherwise a cudaLaunchHostFunc.""" + bs, out = pending + stream = torch.cuda.current_stream().cuda_stream + if self._flag_sync: + # Front-end WAIT(done >= 1): blocks this stream's later nodes without # occupying an SM, so GPU utilization stays truthful during the CPU window. - self._cpu_moe.memop_sync( - torch.cuda.current_stream().cuda_stream, self._done.data_ptr(), slot, - ) + self._ext.flag_wait(stream) else: - stream = torch.cuda.current_stream().cuda_stream - self._ext.sync_with_cuda_stream(stream, task) - io = self._io[bs] - out.copy_(io["y"], non_blocking=True) + self._ext.sync_with_cuda_stream(stream) + out.copy_(self._io["y"][:bs], non_blocking=True) return out - def _watchdog_tick(self, suspects: dict) -> None: + def _watchdog_tick(self, suspect: dict) -> None: """One watchdog sampling round (called every 2 s by ``_watchdog_main``). - A slot is only declared dead when THREE things hold across >=10 s: its doorbell - is still pending (ready==1 && done==0), it was already pending when first - suspected, and the coordinator has served NOTHING on it since (flag_served_count - unchanged). The served-count criterion kills the false-positive window: two - point samples can land on the same slot's (independent, us-scale) pending - windows under heavy external load, but a coordinator that made progress in - between is alive by definition. ``suspects`` maps slot -> (first_seen, - served_at_first_sight) and persists across ticks.""" - stuck = (self._ready == 1) & (self._done == 0) - if not bool(stuck.any()): - suspects.clear() + The doorbell is only declared dead when THREE things hold across >=10 s: it is + still pending (ready set, done==0), it was already pending when first suspected, and + the coordinator has served NOTHING since (flag_served unchanged). The served-count + criterion kills the false-positive window: two point samples can land on + (independent, us-scale) pending windows under heavy external load, but a + coordinator that made progress in between is alive by definition. ``suspect`` + holds (first_seen, served_at_first_sight) and persists across ticks.""" + if not self._ext.flag_pending(): + suspect.clear() return now = time.monotonic() - pending = set(stuck.nonzero().flatten().tolist()) - for slot in list(suspects): - if slot not in pending: - del suspects[slot] - dead = [] - for slot in pending: - served = self._ext.flag_served_count(slot) - first_seen, served_then = suspects.get(slot, (None, None)) - if first_seen is None or served != served_then: - suspects[slot] = (now, served) # new suspect, or alive-but-loaded: rearm - continue - if now - first_seen >= 10.0: - dead.append(slot) - if not dead: + served = self._ext.flag_served() + if not suspect or suspect["served"] != served: + suspect.update(first_seen=now, served=served) # new suspect, or alive-but-loaded: rearm + return + if now - suspect["first_seen"] < 10.0: return logger.error( - f"cpu-moe flag watchdog: slots {dead} unanswered for >10s with no coordinator " - "progress (wedged/dead); poisoning done[] and failing the next step" + "cpu-moe flag watchdog: doorbell unanswered for >10s with no coordinator " + "progress (wedged/dead); poisoning done and failing the next step" ) - for i in dead: - self._err[i] = 1 - for i in dead: - self._done[i] = 1 # after err: unblock the stream into a checked failure - suspects.pop(i, None) + self._unhealthy = True + self._ext.poison_flag() # after the mark: unblock the stream into a checked failure + suspect.clear() def raise_if_unhealthy(self) -> None: """Raise if the flag watchdog fired (a doorbell stayed unanswered because the coordinator never responded). Called by the engine once per forward -- a single - pinned read -- so a dead coordinator surfaces as a loud error on the next step + flag read -- so a dead coordinator surfaces as a loud error on the next step instead of silently shipping stale expert outputs.""" - if self._err is not None and bool((self._err != 0).any()): + if self._unhealthy: raise RuntimeError( "CPU MoE flag-handshake watchdog fired: a decode step's doorbell was " "never answered by the coordinator thread (its outputs cannot be " @@ -728,13 +601,13 @@ def raise_if_unhealthy(self) -> None: def _watchdog_main(executor_ref) -> None: """Watchdog daemon body: weakref-deref per tick so the thread never keeps a dead executor alive (see the start site in ``CpuMoeExecutor.__init__``).""" - suspects: dict = {} + suspect: dict = {} while True: time.sleep(2.0) executor = executor_ref() - if executor is None or executor._watchdog_stop or executor._ready is None: + if executor is None or executor._watchdog_stop: return try: - executor._watchdog_tick(suspects) + executor._watchdog_tick(suspect) finally: del executor # drop the strong ref before the next sleep diff --git a/setup.py b/setup.py index b8e4693c89..5ca5574ff3 100644 --- a/setup.py +++ b/setup.py @@ -112,6 +112,11 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: name="freetoken.kernel._cpu_moe", sources=[ "python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp", + f"python/freetoken/kernel/csrc/cpu_moe/gpu_sync_{'hip' if IS_ROCM else 'cuda'}.cpp", + *sorted( + str(p.relative_to(ROOT)) + for p in (ROOT / "python/freetoken/kernel/csrc/cpu_moe/kernels").glob("*.cpp") + ), ], include_dirs=[KERNEL_INCLUDE, *runtime_include_dirs], library_dirs=runtime_library_dirs, diff --git a/tests/moe/test_cpu_moe.py b/tests/moe/test_cpu_moe.py index 547c94c60a..12bdc47bec 100644 --- a/tests/moe/test_cpu_moe.py +++ b/tests/moe/test_cpu_moe.py @@ -467,7 +467,7 @@ def test_cpu_moe_decode_cuda_graph_replay(): ids = torch.randint(0, E, (bs, top_k), device=dev, dtype=torch.int32) w = torch.rand(bs, top_k, device=dev, dtype=torch.float32) - # Eager forward: also materializes the per-bs pinned buffers + per-(layer,bs) task. + # Eager forward. out_eager = ex.decode(layer, hidden, w, ids) torch.cuda.synchronize() ref_eager = _reference(cache, layer, hidden, ids.cpu(), w.cpu()) @@ -500,11 +500,9 @@ def test_cpu_moe_decode_cuda_graph_replay(): # must actually have exercised it -- guard against a silent fallback to the # host-func path making this test vacuous for the capture-embedded memops. if ex._flag_sync: - assert (layer, bs) in ex._flag_slots, "flag slot expected for the decode task" - slot = ex._flag_slots[(layer, bs)] - assert ex._ext.flag_served_count(slot) >= 4, "1 eager + 3 replay dispatches expected" - assert int(ex._done[slot]) == 1 and int(ex._ready[slot]) == 0, "handshake at rest" - assert int(ex._err.sum()) == 0, "watchdog must not fire in normal operation" + assert ex._ext.flag_served() >= 4, "1 eager + 3 replay dispatches expected" + assert not ex._ext.flag_pending(), "handshake at rest" + assert not ex._unhealthy, "watchdog must not fire in normal operation" ex.raise_if_unhealthy() print("cpu moe cuda graph replay OK") diff --git a/tests/moe/test_cpu_moe_fp8_block.py b/tests/moe/test_cpu_moe_fp8_block.py index 06f954d8d9..baecb86f9a 100644 --- a/tests/moe/test_cpu_moe_fp8_block.py +++ b/tests/moe/test_cpu_moe_fp8_block.py @@ -74,11 +74,11 @@ def test_cpu_fp8_block_accepts_qwen38_padded_scale_bank(): device=torch.device("cpu"), ) hidden = torch.randn(1, H, dtype=torch.bfloat16) - io = ex._io_for(1) + io = ex._io io["x"].copy_(hidden) io["ids"].zero_() io["w"].fill_(1.0) - ex._ext.run_task(ex._task_for(0, 1)) + ex._ext.run(0, 1) cpu_out = io["y"][0].float().to(torch.float64) banks = cache.bank_sources gate_up = _reference_fp8_block_gemv(banks["gate_up"][0], banks["gate_up_scale"][0], hidden[0]) @@ -140,11 +140,8 @@ def _reference_fp8_block_gemv(weights, scales, x): @pytest.mark.parametrize("codes", ["finite", "subnormal"]) -@pytest.mark.parametrize( - "isa, expected_isa", - [("scalar", "scalar"), ("avx2", "avx2"), ("avx512", "avx512f"), ("avx512bf16", "avx512f")], -) -def test_cpu_fp8_block_matches_float64_reference(monkeypatch, isa, expected_isa, codes): +@pytest.mark.parametrize("isa", ["scalar", "avx2", "avx512"]) +def test_cpu_fp8_block_matches_float64_reference(monkeypatch, isa, codes): from freetoken.moe.cpu_executor import CpuMoeExecutor H, I = 259, 257 # ragged K for both GEMVs; 257 down rows exercise many reductions @@ -159,17 +156,16 @@ def test_cpu_fp8_block_matches_float64_reference(monkeypatch, isa, expected_isa, max_tokens=1, device=torch.device("cpu"), ) - tiers = {"scalar": 0, "avx2": 1, "avx512": 2, "avx512bf16": 3} + tiers = {"scalar": 0, "avx2": 1, "avx512": 2} if tiers[isa] > ("scalar", "avx2", "avx512f", "avx512bf16").index(ex.isa): pytest.skip(f"requested {isa}, host/build selected {ex.isa}") - assert ex._ext.fp8_block_isa_name() == expected_isa hidden = torch.randn(1, H, dtype=torch.bfloat16) - io = ex._io_for(1) + io = ex._io io["x"].copy_(hidden) io["ids"].zero_() io["w"].fill_(1.0) - ex._ext.run_task(ex._task_for(0, 1)) + ex._ext.run(0, 1) cpu_out = io["y"][0].float().to(torch.float64) banks = cache.bank_sources @@ -233,9 +229,8 @@ def test_cpu_fp8_block_matches_gpu_decode_kernel(bs, isa, monkeypatch): max_tokens=bs, device=dev, ) - got = ex._ext.fp8_block_isa_name() - if got != _ISA_TIERS[isa]: - pytest.skip(f"asked for {isa}, kernel selected {got}: not supported here") + if ex.isa != _ISA_TIERS[isa]: + pytest.skip(f"asked for {isa}, host/build selected {ex.isa}: not supported here") hidden = torch.randn(bs, H, device=dev, dtype=torch.bfloat16) ids = torch.stack([torch.randperm(E, device=dev)[:top_k] for _ in range(bs)]).to(torch.int32) @@ -254,7 +249,7 @@ def test_cpu_fp8_block_matches_gpu_decode_kernel(bs, isa, monkeypatch): ).float() rel = (cpu_out - gpu_out).abs().max() / (gpu_out.abs().max() + 1e-6) - assert rel < 2e-2, f"bs={bs} isa={got} rel err {rel.item()}" + assert rel < 2e-2, f"bs={bs} isa={ex.isa} rel err {rel.item()}" def test_fp16_bits_carry_every_e4m3_byte(): diff --git a/tests/moe/test_cpu_moe_q4_0.py b/tests/moe/test_cpu_moe_q4_0.py index b5ec2f09b4..7babb5d4ce 100644 --- a/tests/moe/test_cpu_moe_q4_0.py +++ b/tests/moe/test_cpu_moe_q4_0.py @@ -1,7 +1,7 @@ """CPU MoE executor -- native GGUF Q4_0 experts (``--moe-backend cpu``/``hybrid`` for the gemma4 GGUF checkpoint). -The CPU W4A16 GEMV (``q4_0_dot`` in csrc/cpu_moe/cpu_moe_ext.cpp) reads the *same* +The CPU W4A16 GEMV (``q4_0_dot`` in csrc/cpu_moe/kernels/q4_0.cpp) reads the *same* packed Q4_0 banks the GPU offload path streams and dequantizes weights inside the K-loop. We check it against the reference dequant (models/gguf/dequant.py) + the production bf16 GPU decode kernel on byte-identical banks: both are W4A16, so the diff --git a/tests/moe/test_dsfp4_prequant.py b/tests/moe/test_dsfp4_prequant.py index 11519502e2..8187402d3f 100644 --- a/tests/moe/test_dsfp4_prequant.py +++ b/tests/moe/test_dsfp4_prequant.py @@ -17,8 +17,9 @@ @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") -def test_dsfp4_gpu_prequant_bit_parity(): +def test_dsfp4_gpu_prequant_bit_parity(monkeypatch): from freetoken.kernel.pinned import alloc_pinned_tensor + from freetoken.moe import cpu_executor from freetoken.moe.cpu_executor import CpuMoeExecutor torch.manual_seed(0) @@ -44,14 +45,13 @@ def test_dsfp4_gpu_prequant_bit_parity(): wts = torch.rand(2, TOPK, device=dev) def run(cpu_roundtrip: bool) -> torch.Tensor: + # control: the original CPU-side scalar round-trip + monkeypatch.setattr(cpu_executor, "_DSFP4_GPU_PREQUANT", not cpu_roundtrip) ex = CpuMoeExecutor( cache, top_k=TOPK, activation="silu", apply_router_weight_on_input=False, num_threads=4, max_tokens=2, device=dev, swiglu_limit=10.0, act_block=128, ) - assert ex._gpu_prequant, "ds_fp4 on CUDA should enable GPU prequant" - if cpu_roundtrip: # control: the original CPU-side scalar round-trip - ex._gpu_prequant = False - ex._ext.set_input_prequant(False) + assert ex._gpu_prequant != cpu_roundtrip, "ds_fp4 on CUDA should enable GPU prequant" out = ex.decode(0, x, wts, ids.clone()).clone() torch.cuda.synchronize() del ex From d975292330e27bd2a7480306aee488e52fb0e381 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Tue, 13 Oct 2026 23:59:13 +0000 Subject: [PATCH 2/4] perf(cpu-moe): skip the worker pool for tasks without routed experts --- .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 18 ++++++++++++++++-- 1 file changed, 16 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 40cbd6796c..a8a6868f67 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -217,6 +217,14 @@ struct CpuMoeExecutor { std::string isa_name() const { return fmt->isa_name(); } + // Whether any route of the task names an expert; do_pass1 and do_pass2 skip the others. + bool has_routes(const MoeTask& t) const { + const int64_t n = static_cast(t.num_tokens) * params.top_k; + for (int64_t i = 0; i < n; ++i) + if (io_ids[i] >= 0 && io_ids[i] < num_experts) return true; + return false; + } + void barrier(int& local_sense) { local_sense ^= 1; if (bar_count.fetch_add(1) + 1 == num_threads) { @@ -430,8 +438,14 @@ struct CpuMoeExecutor { while (!coord_stop.load(std::memory_order_relaxed)) { const uint64_t word = gpu_sync->cpu_poll(); if (word != 0) { - submit(from_word(word)); - sync(coord_pinned && spin_wait.load(std::memory_order_relaxed)); + const MoeTask t = from_word(word); + if (has_routes(t)) { + submit(t); + sync(coord_pinned && spin_wait.load(std::memory_order_relaxed)); + } else { + // What the pool would write for a task without routes, minus its fork and join. + std::memset(io_y, 0, sizeof(bf16_t) * t.num_tokens * params.H); + } // Release: the workers' y stores are visible before the GPU sees done. gpu_sync->cpu_done(); flag_served.fetch_add(1, std::memory_order_relaxed); From 886f75234d461b96531f78f22426c0880d89c3b1 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Wed, 14 Oct 2026 00:58:57 +0000 Subject: [PATCH 3/4] perf(cpu-moe): publish and combine the cpu handshake with kernels --- .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 19 ++ .../freetoken/kernel/csrc/cpu_moe/gpu_sync.h | 6 + .../kernel/csrc/cpu_moe/gpu_sync_cuda.cpp | 10 + .../kernel/csrc/cpu_moe/gpu_sync_hip.cpp | 4 + .../kernel/csrc/cpu_moe/gpu_sync_io.cu | 199 ++++++++++++++++++ .../kernel/csrc/cpu_moe/gpu_sync_io.h | 39 ++++ .../csrc/include/freetoken/hip_compat.h | 12 ++ python/freetoken/layers/moe.py | 5 +- python/freetoken/moe/cpu_executor.py | 43 +++- setup.py | 12 +- 10 files changed, 337 insertions(+), 12 deletions(-) create mode 100644 python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu create mode 100644 python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h 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 a8a6868f67..c8b4d1b465 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -37,6 +37,7 @@ #include "format.h" #include "gpu_sync.h" +#include "gpu_sync_io.h" namespace freetoken::cpu_moe { @@ -158,6 +159,7 @@ struct CpuMoeExecutor { // The GPU rings the doorbell with a task word at submit; this coordinator thread busy-polls it, runs the task on the worker pool, and completes it, which the GPU waits on at sync. // This removes the ~2x30-50us host-func dispatch round trips per MoE layer per decode step that otherwise idle the GPU (~6 ms/step on a 75-layer model). std::unique_ptr gpu_sync; + std::unique_ptr kernel_io; // replaces the copies around the doorbell once enable_kernel_io() runs std::thread coord_thread; std::atomic coord_stop{false}; std::atomic flag_served{0}; // completed doorbells (tests, watchdog) @@ -411,6 +413,17 @@ struct CpuMoeExecutor { gpu_sync->gpu_submit(stream, to_word(task(layer_id, num_tokens))); } void flag_wait(uintptr_t stream) { gpu_sync->gpu_wait(stream); } + void enable_kernel_io() { + kernel_io = std::make_unique(gpu_sync->ready_word(), gpu_sync->done_word(), io_x, io_ids, io_w, io_y, params.H, + params.top_k); + } + void flag_publish(uintptr_t stream, uintptr_t x, uintptr_t ids, uintptr_t w, int layer_id, int num_tokens) { + kernel_io->publish(stream, reinterpret_cast(x), reinterpret_cast(ids), + reinterpret_cast(w), num_tokens, to_word(task(layer_id, num_tokens))); + } + void flag_combine(uintptr_t stream, uintptr_t out, int num_tokens, bool accumulate) { + kernel_io->combine(stream, reinterpret_cast(out), num_tokens, accumulate); + } bool flag_pending() const { return gpu_sync && gpu_sync->pending(); } // Complete a doorbell the coordinator never answered so the stream can make progress. void poison_flag() { @@ -520,6 +533,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("num_tokens")) .def("flag_wait", &CpuMoeExecutor::flag_wait, py::arg("stream")) .def("flag_pending", &CpuMoeExecutor::flag_pending) + .def("enable_kernel_io", &CpuMoeExecutor::enable_kernel_io) + .def("flag_publish", &CpuMoeExecutor::flag_publish, py::arg("stream"), py::arg("x"), py::arg("ids"), py::arg("w"), + py::arg("layer_id"), py::arg("num_tokens")) + .def("flag_combine", &CpuMoeExecutor::flag_combine, py::arg("stream"), py::arg("out"), py::arg("num_tokens"), + py::arg("accumulate")) .def("poison_flag", &CpuMoeExecutor::poison_flag) .def("isa_name", &CpuMoeExecutor::isa_name) .def("set_spin_wait", @@ -527,6 +545,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("enabled")) .def("get_spin_wait", [](const CpuMoeExecutor& e) { return e.spin_wait.load(); }); m.def("memops_probe", &GpuSync::probe, py::arg("stream")); + m.def("kernel_io_probe", &freetoken::cpu_moe::KernelIo::probe, py::arg("stream")); m.def("host_func_capture_safe", &GpuSync::host_func_capture_safe); // ABI capability marker: the highest ActKind this build implements in the // GENERIC epilogue. CpuMoeExecutor.__init__ probes it before requesting an act diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h index bf58874168..829a8f34c3 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h @@ -5,6 +5,8 @@ // 2. wait, CPU -> GPU: gpu_wait(stream) holds later work of the stream until cpu_done(), and that work sees every CPU write made before cpu_done(). // task is a nonzero word, opaque here; the executor encodes the layer and token count in it, as a constant baked into each captured graph. // One task is in flight at a time: decode runs on one stream and every layer waits before the next one signals. +// KernelIo (gpu_sync_io.h) may replace the copies around submit and wait: its kernel copies the inputs into the IO rows and does gpu_submit's writes itself. +// It skips a step without a CPU route: done = 1 and no ready, so cpu_poll() never sees the step and gpu_wait() passes at once. #pragma once #include @@ -64,6 +66,10 @@ class GpuSync { // Submitted but not polled yet (watchdog only). bool pending() const; + // Device addresses of the two words, for KernelIo. + uintptr_t ready_word() const; + uintptr_t done_word() const; + private: struct Impl; std::unique_ptr impl_; diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp index 0b73a91d61..a2a089bf67 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp @@ -89,6 +89,7 @@ static void cumemop_check(int rc, const char* what) { // ready and done live on separate cache lines of one mapped-pinned block; with UVA the GPU memops and the coordinator use the same addresses. struct GpuSync::Impl { void* block = nullptr; + void* device_block = nullptr; volatile int64_t* ready = nullptr; volatile int64_t* done = nullptr; }; @@ -101,6 +102,11 @@ GpuSync::GpuSync() : impl_(std::make_unique()) { std::memset(impl_->block, 0, 128); impl_->ready = static_cast(impl_->block); impl_->done = impl_->ready + 8; + const cudaError_t dev_rc = cudaHostGetDevicePointer(&impl_->device_block, impl_->block, 0); + if (dev_rc != cudaSuccess) { + cudaFreeHost(impl_->block); + throw std::runtime_error(std::string("CPU MoE flag device address: ") + cudaGetErrorString(dev_rc)); + } } GpuSync::~GpuSync() { cudaFreeHost(impl_->block); } @@ -153,4 +159,8 @@ void GpuSync::cpu_done() { flag_store_release(impl_->done, 1); } bool GpuSync::pending() const { return flag_load_acquire(impl_->ready) != 0 && flag_load_acquire(impl_->done) == 0; } +uintptr_t GpuSync::ready_word() const { return reinterpret_cast(impl_->device_block); } + +uintptr_t GpuSync::done_word() const { return reinterpret_cast(impl_->device_block) + 8 * sizeof(int64_t); } + } // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp index 4baf76e768..eb9e7dc662 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp @@ -257,4 +257,8 @@ bool GpuSync::pending() const { return flag_load_acquire(as_flag(impl_->ready)) != 0 && flag_load_acquire(as_flag(impl_->done)) == 0; } +uintptr_t GpuSync::ready_word() const { return reinterpret_cast(impl_->ready); } + +uintptr_t GpuSync::done_word() const { return reinterpret_cast(impl_->done); } + } // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu new file mode 100644 index 0000000000..b41961f0bf --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu @@ -0,0 +1,199 @@ +#include +#include +#include + +#include + +#include "gpu_sync_io.h" + +namespace freetoken::cpu_moe { +namespace { + +__device__ inline float bf16_bits_to_float(uint16_t b) { return __uint_as_float(static_cast(b) << 16); } + +// Round to nearest even, the same rounding as a torch bf16 add. +__device__ inline uint16_t float_to_bf16_bits(float f) { + const uint32_t u = __float_as_uint(f); + if ((u & 0x7fffffffu) > 0x7f800000u) return 0x7fc0; + return static_cast((u + 0x7fffu + ((u >> 16) & 1u)) >> 16); +} + +// One block. Copies x/ids/w into the IO rows, then done = 0 and ready = task. +// A step without a CPU route writes done = 1 and rings nothing: the coordinator stays asleep and the GPU wait passes at once. +__global__ void publish_kernel(const uint16_t* __restrict__ x, const int32_t* __restrict__ ids, const float* __restrict__ w, + int64_t n_x, int64_t n_ids, uint16_t* io_x, int32_t* io_ids, float* io_w, + volatile int64_t* ready, volatile int64_t* done, int64_t task, int32_t* has_cpu) { + __shared__ int any; + if (threadIdx.x == 0) any = 0; + __syncthreads(); + for (int64_t i = threadIdx.x; i < n_ids; i += blockDim.x) + if (ids[i] >= 0) any = 1; + __syncthreads(); + if (!any) { + if (threadIdx.x == 0) { + *has_cpu = 0; + *done = 1; + __threadfence_system(); + } + return; + } + const bool vec = (reinterpret_cast(x) | reinterpret_cast(io_x)) % 16 == 0 && n_x % 8 == 0; + if (vec) { + for (int64_t i = threadIdx.x; i < n_x / 8; i += blockDim.x) + reinterpret_cast(io_x)[i] = reinterpret_cast(x)[i]; + } else { + for (int64_t i = threadIdx.x; i < n_x; i += blockDim.x) io_x[i] = x[i]; + } + for (int64_t i = threadIdx.x; i < n_ids; i += blockDim.x) { + io_ids[i] = ids[i]; + io_w[i] = w[i]; + } + __threadfence_system(); + __syncthreads(); + if (threadIdx.x == 0) { + *has_cpu = 1; + *done = 0; + __threadfence_system(); + *ready = task; + __threadfence_system(); + } +} + +// out = CPU rows, or out += CPU rows with accumulate; a step publish skipped contributes zeros. +// io_y is volatile: the same host rows hold a new step on every call. +__global__ void combine_kernel(uint16_t* out, const volatile uint16_t* io_y, const int32_t* has_cpu, int64_t n, bool accumulate) { + const bool on = *has_cpu != 0; + if (accumulate && !on) return; + for (int64_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; i < n; i += static_cast(gridDim.x) * blockDim.x) { + float v = accumulate ? bf16_bits_to_float(out[i]) : 0.0f; + if (on) v += bf16_bits_to_float(io_y[i]); + out[i] = float_to_bf16_bits(v); + } +} + +cudaError_t launch_publish(cudaStream_t stream, const uint16_t* x, const int32_t* ids, const float* w, int64_t n_x, int64_t n_ids, + uint16_t* io_x, int32_t* io_ids, float* io_w, volatile int64_t* ready, volatile int64_t* done, int64_t task, + int32_t* has_cpu) { + publish_kernel<<<1, 256, 0, stream>>>(x, ids, w, n_x, n_ids, io_x, io_ids, io_w, ready, done, task, has_cpu); + return cudaGetLastError(); +} + +cudaError_t launch_combine(cudaStream_t stream, uint16_t* out, const uint16_t* io_y, const int32_t* has_cpu, int64_t n, bool accumulate) { + if (n == 0) return cudaSuccess; + const auto blocks = static_cast(std::min((n + 255) / 256, 1024)); + combine_kernel<<>>(out, io_y, has_cpu, n, accumulate); + return cudaGetLastError(); +} + +void check(cudaError_t rc, const char* what) { + if (rc != cudaSuccess) throw std::runtime_error(std::string("CPU MoE kernel IO: ") + what + ": " + cudaGetErrorString(rc)); +} + +template +T* device_address(const void* host) { + void* dev = nullptr; + check(cudaHostGetDevicePointer(&dev, const_cast(host), 0), "device address of the IO rows"); + return static_cast(dev); +} + +} // namespace + +KernelIo::KernelIo(uintptr_t ready, uintptr_t done, const void* x, const int32_t* ids, const float* w, const void* y, int hidden, + int top_k) + : ready_(reinterpret_cast(ready)), + done_(reinterpret_cast(done)), + x_(device_address(x)), + ids_(device_address(ids)), + w_(device_address(w)), + y_(device_address(y)), + hidden_(hidden), + top_k_(top_k) { + check(cudaMalloc(reinterpret_cast(&has_cpu_), sizeof(int32_t)), "has_cpu allocation"); + const cudaError_t rc = cudaMemset(has_cpu_, 0, sizeof(int32_t)); + if (rc != cudaSuccess) { + (void)cudaFree(has_cpu_); + check(rc, "has_cpu reset"); + } +} + +KernelIo::~KernelIo() { (void)cudaFree(has_cpu_); } + +void KernelIo::publish(uintptr_t stream, const void* x, const int32_t* ids, const float* w, int num_tokens, uint64_t task) { + check(launch_publish(reinterpret_cast(stream), static_cast(x), ids, w, + static_cast(num_tokens) * hidden_, static_cast(num_tokens) * top_k_, x_, ids_, w_, ready_, done_, + static_cast(task), has_cpu_), + "publish launch"); +} + +void KernelIo::combine(uintptr_t stream, void* out, int num_tokens, bool accumulate) { + check(launch_combine(reinterpret_cast(stream), static_cast(out), y_, has_cpu_, + static_cast(num_tokens) * hidden_, accumulate), + "combine launch"); +} + +// A ringing step, its combine and a skipped step on scratch mapped memory, checking every word they write. +std::string KernelIo::probe(uintptr_t stream_handle) { + struct alignas(16) Area { + int64_t ready, done; + int32_t has_cpu, pad[3]; + uint16_t x[8], io_x[8], io_y[8], out[8]; + int32_t ids[2], io_ids[2]; + float w[2], io_w[2]; + }; + const auto stream = reinterpret_cast(stream_handle); + void* host = nullptr; + void* dev = nullptr; + if (cudaHostAlloc(&host, sizeof(Area), cudaHostAllocPortable | cudaHostAllocMapped) != cudaSuccess) + return "pinned scratch allocation failed"; + std::string reason; + auto ok = [&](cudaError_t rc, const char* what) { + if (rc != cudaSuccess && reason.empty()) reason = std::string(what) + ": " + cudaGetErrorString(rc); + return rc == cudaSuccess; + }; + auto fail = [&](const char* why) { + if (reason.empty()) reason = why; + }; + if (ok(cudaHostGetDevicePointer(&dev, host, 0), "device address of the scratch")) { + auto* h = static_cast(host); + auto* d = static_cast(dev); + std::memset(h, 0, sizeof(Area)); + for (int i = 0; i < 8; ++i) { + h->x[i] = static_cast(0x3f80 + i); + h->io_y[i] = 0x3f80; // 1.0 + h->out[i] = 0x4000; // 2.0 + } + h->ids[0] = 3; + h->ids[1] = -1; + h->w[0] = 0.25f; + h->w[1] = 0.75f; + auto publish = [&] { + return launch_publish(stream, d->x, d->ids, d->w, 8, 2, d->io_x, d->io_ids, d->io_w, &d->ready, &d->done, 5, &d->has_cpu); + }; + + if (ok(publish(), "publish launch") && ok(cudaStreamSynchronize(stream), "publish")) { + if (h->ready != 5 || h->done != 0 || h->has_cpu != 1) fail("publish wrote wrong handshake words"); + if (std::memcmp(h->io_x, h->x, sizeof(h->x)) || std::memcmp(h->io_ids, h->ids, sizeof(h->ids)) || + std::memcmp(h->io_w, h->w, sizeof(h->w))) + fail("publish wrote wrong IO rows"); + } + if (reason.empty() && ok(launch_combine(stream, d->out, d->io_y, &d->has_cpu, 8, true), "combine launch") && + ok(cudaStreamSynchronize(stream), "combine")) { + for (int i = 0; i < 8; ++i) + if (h->out[i] != 0x4040) fail("combine returned a wrong sum"); // 3.0 + } + if (reason.empty()) { + h->ready = 0; + h->ids[0] = -1; + if (ok(publish(), "skip launch") && ok(launch_combine(stream, d->out, d->io_y, &d->has_cpu, 8, false), "skip combine launch") && + ok(cudaStreamSynchronize(stream), "skipped step")) { + bool zeros = true; + for (int i = 0; i < 8; ++i) zeros = zeros && h->out[i] == 0; + if (h->ready != 0 || h->done != 1 || h->has_cpu != 0 || !zeros) fail("a skipped step wrote wrong words"); + } + } + } + (void)cudaFreeHost(host); + return reason; +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h new file mode 100644 index 0000000000..6eacb3936b --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h @@ -0,0 +1,39 @@ +// Kernel IO for the GPU <-> CPU handshake (gpu_sync.h): kernels on the executor's mapped pinned IO rows in place of the copies around submit and wait. +// One source for both runtimes: nvcc builds gpu_sync_io.cu on CUDA, hipcc on ROCm through freetoken/hip_compat.h. +#pragma once + +#include +#include + +namespace freetoken::cpu_moe { + +class KernelIo { + public: + // ready, done: device addresses of GpuSync's words. x, ids, w, y: host addresses of the pinned IO rows. Throws if setup fails. + KernelIo(uintptr_t ready, uintptr_t done, const void* x, const int32_t* ids, const float* w, const void* y, int hidden, int top_k); + ~KernelIo(); + KernelIo(const KernelIo&) = delete; + KernelIo& operator=(const KernelIo&) = delete; + + // Copies the step's device inputs into the IO rows and does gpu_submit's writes, or skips the step when no route goes to the CPU. + // x [n, hidden] bf16, ids [n, top_k] int32 (< 0 = not a CPU route), w [n, top_k] fp32. + void publish(uintptr_t stream, const void* x, const int32_t* ids, const float* w, int num_tokens, uint64_t task); + // Follows gpu_wait(): writes the CPU rows to out [n, hidden] bf16, or adds them to it; a skipped step gives zeros. + void combine(uintptr_t stream, void* out, int num_tokens, bool accumulate); + + // Whether both kernels run and write the right words on this device: "" or the reason. + static std::string probe(uintptr_t stream); + + private: + volatile int64_t* ready_; + volatile int64_t* done_; + uint16_t* x_; + int32_t* ids_; + float* w_; + const uint16_t* y_; + int hidden_; + int top_k_; + int32_t* has_cpu_ = nullptr; // device word: whether the last published step has CPU routes +}; + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h index 1898b41189..8abcd1a098 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h +++ b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h @@ -44,6 +44,18 @@ #define cudaGetLastError hipGetLastError #endif +#ifndef cudaMalloc +#define cudaMalloc hipMalloc +#endif + +#ifndef cudaFree +#define cudaFree hipFree +#endif + +#ifndef cudaMemset +#define cudaMemset hipMemset +#endif + #ifndef cudaMallocHost #define cudaMallocHost hipHostMalloc #endif diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 4918a032fe..82f09a9bb8 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -346,8 +346,9 @@ def _decode_hybrid( alphas=cache.alphas_for_slots(self.layer_id), is_prefill=False, ) - cpu_routed = cpu_routed_early if not _HYBRID_OVERLAP else executor.decode_sync(pending) - return gpu_routed + cpu_routed + if not _HYBRID_OVERLAP: + return gpu_routed + cpu_routed_early + return executor.decode_sync(pending, add_to=gpu_routed) def _prefill_routed( self, diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index 079c1f43d3..bb9ee56385 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -8,6 +8,8 @@ -> submit (doorbell naming (layer, bs), or a cudaLaunchHostFunc host node) -> sync (wait for the CPU pool to finish) -> H2D (pinned expert output -> GPU) + +With the doorbell on CUDA, a publish kernel does the D2H and the doorbell, and skips the step when no route goes to the CPU; a combine kernel does the H2D (csrc/cpu_moe/gpu_sync_io.h). """ from __future__ import annotations @@ -302,8 +304,17 @@ def __init__( # stream and raise_if_unhealthy() turns the step into a loud error instead of silent # stale output (flag_sync itself was decided above, before thread sizing). self._unhealthy = False + self._kernel_io = False if self._flag_sync: self._ext.start_flag_coordinator(self._coord_core) + # ROCm keeps the copies until the kernel IO has been checked on AMD GPUs. + if torch.version.hip is None: + reason = _cpu_moe.kernel_io_probe(torch.cuda.current_stream().cuda_stream) + if reason: + logger.warning_rank0(f"cpu-moe kernel IO unavailable ({reason}); using D2H/H2D copies around the doorbell") + else: + self._ext.enable_kernel_io() + self._kernel_io = True self._watchdog_stop = False # The thread target holds a WEAKREF and re-derefs it each tick: a bound # method would strong-reference the executor forever (the loop never ends @@ -326,7 +337,8 @@ def __init__( f"{core_ids[0]}..{core_ids[-1]}) isa={self.isa} fmt={fmt} " f"H={self.H} I={self.I} experts={self.num_experts} layers={self.num_layers} " f"top_k={self.top_k} act={activation} max_tokens={self.max_tokens} " - f"spin_wait={self.spin_wait}" + f"spin_wait={self.spin_wait} " + f"handshake={'kernel' if self._kernel_io else 'memop' if self._flag_sync else 'host-func'}" ) def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor: @@ -498,7 +510,7 @@ def decode_submit( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> tuple: - """Issue the D2H copies + the CPU-pool submit host node, then return without + """Issue the D2H copies + the CPU-pool submit host node (or the publish kernel), then return without waiting. Lets a caller (the hybrid backend) enqueue GPU work between this and :meth:`decode_sync` so the CPU compute overlaps the GPU GEMM / PCIe fetch. @@ -524,6 +536,15 @@ def decode_submit( hidden_states = act_quant_fp8_roundtrip(hidden_states, block=self._act_block) + stream = torch.cuda.current_stream().cuda_stream + if self._kernel_io and hidden_states.dtype == torch.bfloat16: + x = hidden_states.contiguous() + ids = topk_ids.to(torch.int32).contiguous() + w = topk_weights.to(torch.float32).contiguous() + assert x.shape == (bs, self.H) and ids.shape == w.shape == (bs, self.top_k), (x.shape, ids.shape, w.shape) + self._ext.flag_publish(stream, x.data_ptr(), ids.data_ptr(), w.data_ptr(), layer_id, bs) + return (bs, torch.empty_like(x), True) + # D2H: ship this step's activations + routing to pinned host memory. io = self._io io["x"][:bs].copy_(hidden_states, non_blocking=True) @@ -531,20 +552,20 @@ def decode_submit( io["w"][:bs].copy_(topk_weights.to(torch.float32), non_blocking=True) out = torch.empty_like(hidden_states) - stream = torch.cuda.current_stream().cuda_stream if self._flag_sync: # Front-end memops: reset done, then ring the doorbell with (layer, bs). # No kernel launched; no host-func round trip. self._ext.flag_signal(stream, layer_id, bs) else: self._ext.submit_with_cuda_stream(stream, layer_id, bs) - return (bs, out) + return (bs, out, False) - def decode_sync(self, pending: tuple) -> torch.Tensor: + def decode_sync(self, pending: tuple, add_to: torch.Tensor | None = None) -> torch.Tensor: """Issue the CPU-pool sync + the H2D result copy for a prior :meth:`decode_submit`, and return the GPU output tensor. With flag-sync the wait is a front-end stream - memop on done (set by the CPU coordinator); otherwise a cudaLaunchHostFunc.""" - bs, out = pending + memop on done (set by the CPU coordinator); otherwise a cudaLaunchHostFunc. + With ``add_to``, return ``add_to`` + the CPU output; the kernel path adds in place into ``add_to``.""" + bs, out, kernel_io = pending stream = torch.cuda.current_stream().cuda_stream if self._flag_sync: # Front-end WAIT(done >= 1): blocks this stream's later nodes without @@ -552,8 +573,14 @@ def decode_sync(self, pending: tuple) -> torch.Tensor: self._ext.flag_wait(stream) else: self._ext.sync_with_cuda_stream(stream) + if kernel_io: + if add_to is not None and add_to.dtype == torch.bfloat16 and add_to.is_contiguous() and add_to.shape == out.shape: + self._ext.flag_combine(stream, add_to.data_ptr(), bs, True) + return add_to + self._ext.flag_combine(stream, out.data_ptr(), bs, False) + return out if add_to is None else add_to + out out.copy_(self._io["y"][:bs], non_blocking=True) - return out + return out if add_to is None else add_to + out def _watchdog_tick(self, suspect: dict) -> None: """One watchdog sampling round (called every 2 s by ``_watchdog_main``). diff --git a/setup.py b/setup.py index 5ca5574ff3..b968373151 100644 --- a/setup.py +++ b/setup.py @@ -77,15 +77,22 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: if IS_ROCM: runtime_include_dirs, runtime_library_dirs, runtime_lib = _rocm_paths() runtime_link_args = [f"-Wl,-rpath,{runtime_library_dirs[0]}"] - # These extensions contain host code only. BuildExtension supplies the ROCm + # Host sources. BuildExtension supplies the ROCm # platform defines to the C++ compiler; offload architecture flags belong on # HIP device sources and would be rejected by the host compiler here. extra_compile = ["-O3", "-std=c++20"] + # Device sources (.cu, built by hipcc); BuildExtension adds the offload architectures. + device_compile = ["-O3", "-std=c++20"] else: runtime_include_dirs, runtime_library_dirs = _cuda_runtime_paths() runtime_lib = "cudart" runtime_link_args = [] extra_compile = ["-O3", "-std=c++20"] + # Device sources (.cu, built by nvcc): SASS for each major arch, which also runs on the newer minors of that major, + # plus PTX for the oldest arch nvcc 13 supports, which the driver builds for any other GPU at load. + # Naming archs also stops torch from picking them, which fails on a build machine without a GPU. + device_compile = ["-O3", "-std=c++20"] + [f"-gencode=arch=compute_{a},code=sm_{a}" for a in (75, 80, 90, 100, 120)] + device_compile.append("-gencode=arch=compute_75,code=compute_75") _check_toolchain() @@ -113,6 +120,7 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: sources=[ "python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp", f"python/freetoken/kernel/csrc/cpu_moe/gpu_sync_{'hip' if IS_ROCM else 'cuda'}.cpp", + "python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu", *sorted( str(p.relative_to(ROOT)) for p in (ROOT / "python/freetoken/kernel/csrc/cpu_moe/kernels").glob("*.cpp") @@ -121,7 +129,7 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: include_dirs=[KERNEL_INCLUDE, *runtime_include_dirs], library_dirs=runtime_library_dirs, libraries=[runtime_lib], - extra_compile_args=extra_compile + ["-pthread"], + extra_compile_args={"cxx": extra_compile + ["-pthread"], "nvcc": device_compile}, extra_link_args=runtime_link_args, ), # disk-backed row store (PLE / Engram tables); Linux-only until the TableFile/BatchReader seams grow Windows bodies From d334eee5647613b4559a5f8afa322adf9f07cd9f Mon Sep 17 00:00:00 2001 From: zihaomu Date: Tue, 13 Oct 2026 21:42:39 +0800 Subject: [PATCH 4/4] fix(cpu-moe): clean up failed hip capability probes --- .../kernel/csrc/cpu_moe/gpu_sync_hip.cpp | 13 ++- tests/moe/hip_probe_faults.cpp | 75 +++++++++++++ tests/moe/test_cpu_moe.py | 105 ++++++++++++++++++ 3 files changed, 192 insertions(+), 1 deletion(-) create mode 100644 tests/moe/hip_probe_faults.cpp diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp index eb9e7dc662..31ff9be539 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp @@ -28,6 +28,12 @@ void hip_check(hipError_t rc, const char* what) { throw std::runtime_error(std::string("CPU MoE HIP flag sync: ") + what + " failed: " + hipGetErrorString(rc)); } +// Capability failures must not poison eager fallback or hide a caller's earlier error. +struct ProbeErrorScope { + ProbeErrorScope() { hip_check(hipPeekAtLastError(), "pending error before probe"); } + ~ProbeErrorScope() { (void)hipGetLastError(); } +}; + hipError_t alloc_signal(uint64_t** out) { void* storage = nullptr; const hipError_t rc = hipExtMallocWithFlags(&storage, sizeof(uint64_t), hipMallocSignalMemory); @@ -53,6 +59,7 @@ bool capturing(hipStream_t stream) { // Whether this runtime records a native hipStreamWriteValue64 issued under capture as a graph node. bool records_native_memops() { static const bool native = [] { + ProbeErrorScope errors; uint64_t* word = nullptr; hipStream_t stream = nullptr; hipGraph_t graph = nullptr; @@ -106,7 +113,10 @@ hipError_t add_capture_memop_node(hipStream_t stream, hipStreamBatchMemOpParams* hipGraphNode_t node = nullptr; rc = hipGraphAddBatchMemOpNode(&node, graph, dependencies, dependency_count, ¶ms); if (rc != hipSuccess) return rc; - return hipStreamUpdateCaptureDependencies(stream, &node, 1, hipStreamSetCaptureDependencies); + rc = hipStreamUpdateCaptureDependencies(stream, &node, 1, hipStreamSetCaptureDependencies); + // EndCapture rejects a node that was added but never joined to the stream. + if (rc != hipSuccess && node != nullptr) (void)hipGraphDestroyNode(node); + return rc; } hipStreamBatchMemOpParams write_op(uint64_t* address, uint64_t value) { @@ -170,6 +180,7 @@ GpuSync::~GpuSync() { } bool GpuSync::probe(uintptr_t /*stream*/) { + ProbeErrorScope errors; int device = 0; int can_wait = 0; if (hipGetDevice(&device) != hipSuccess || diff --git a/tests/moe/hip_probe_faults.cpp b/tests/moe/hip_probe_faults.cpp new file mode 100644 index 0000000000..4426bf0d14 --- /dev/null +++ b/tests/moe/hip_probe_faults.cpp @@ -0,0 +1,75 @@ +#include +#include + +#include +#include +#include + +namespace { +std::atomic begins{0}, ends{0}, injected{0}, removed_nodes{0}, last_end_rc{0}; + +bool mode(const char* value) { + const char* selected = std::getenv("FREETOKEN_TEST_HIP_PROBE_FAULT"); + return selected != nullptr && std::strcmp(selected, value) == 0; +} + +template +Function real(const char* name) { + return reinterpret_cast(dlsym(RTLD_NEXT, name)); +} +} // namespace + +extern "C" hipError_t hipStreamBeginCapture(hipStream_t stream, hipStreamCaptureMode capture_mode) { + const hipError_t rc = real("hipStreamBeginCapture")(stream, capture_mode); + if (rc == hipSuccess) ++begins; + return rc; +} + +extern "C" hipError_t hipStreamWriteValue64(hipStream_t stream, void* address, uint64_t value, unsigned flags) { + // Force the compatibility path on runtimes that normally record native writes. + if (mode("dependency") && begins == 1) { + ++injected; + return hipSuccess; + } + return real("hipStreamWriteValue64")(stream, address, value, flags); +} + +extern "C" hipError_t hipStreamUpdateCaptureDependencies( + hipStream_t stream, hipGraphNode_t* dependencies, size_t count, unsigned flags) { + if (mode("dependency")) { + ++injected; + return hipErrorNotSupported; + } + return real("hipStreamUpdateCaptureDependencies")( + stream, dependencies, count, flags); +} + +extern "C" hipError_t hipGraphDestroyNode(hipGraphNode_t node) { + const hipError_t rc = real("hipGraphDestroyNode")(node); + if (rc == hipSuccess) ++removed_nodes; + return rc; +} + +extern "C" hipError_t hipStreamEndCapture(hipStream_t stream, hipGraph_t* graph) { + ++ends; + hipError_t rc = real("hipStreamEndCapture")(stream, graph); + if (mode("end-capture") && ends == 2 && rc == hipSuccess) { + // A real invalid API call also sets HIP's thread-local last-error state. + rc = real("hipGraphGetNodes")(*graph, nullptr, nullptr); + ++injected; + } + last_end_rc = static_cast(rc); + return rc; +} + +extern "C" int probe_test_injected() { return injected; } +extern "C" int probe_test_begins() { return begins; } +extern "C" int probe_test_ends() { return ends; } +extern "C" int probe_test_removed_nodes() { return removed_nodes; } +extern "C" int probe_test_last_end_rc() { return last_end_rc; } +extern "C" int probe_test_peek_error() { + return static_cast(real("hipPeekAtLastError")()); +} +extern "C" int probe_test_make_error() { + return static_cast(real("hipGraphGetNodes")(nullptr, nullptr, nullptr)); +} diff --git a/tests/moe/test_cpu_moe.py b/tests/moe/test_cpu_moe.py index 12bdc47bec..38e79b2743 100644 --- a/tests/moe/test_cpu_moe.py +++ b/tests/moe/test_cpu_moe.py @@ -440,6 +440,111 @@ def _reference(cache, layer, hidden, ids, w, act="silu", apply_in=False): return out +@pytest.fixture(scope="module") +def hip_probe_fault_library(tmp_path_factory): + import importlib.util + import os + from pathlib import Path + import shutil + import subprocess + + if not torch.version.hip: + pytest.skip("requires the HIP sync backend") + compiler = shutil.which("c++") + if compiler is None: + pytest.skip("requires a C++ compiler for HIP fault injection") + candidates = [Path(os.getenv("ROCM_HOME", "/opt/rocm")), Path("/opt/rocm")] + spec = importlib.util.find_spec("_rocm_sdk_core") + if spec and spec.submodule_search_locations: + candidates.append(Path(next(iter(spec.submodule_search_locations)))) + include = next((root / "include" for root in candidates if (root / "include/hip/hip_runtime.h").exists()), None) + if include is None: + pytest.skip("requires HIP SDK headers for fault injection") + library = tmp_path_factory.mktemp("hip_probe") / "faults.so" + command = [ + compiler, "-std=c++20", "-shared", "-fPIC", "-D__HIP_PLATFORM_AMD__=1", + f"-I{include}", str(Path(__file__).with_name("hip_probe_faults.cpp")), + "-ldl", "-o", str(library), + ] + result = subprocess.run(command, capture_output=True, text=True, timeout=60) + assert result.returncode == 0, result.stdout + result.stderr + return library + + +@pytest.mark.parametrize("fault", ["dependency", "end-capture", "pending"]) +def test_cpu_moe_hip_probe_error_recovery(hip_probe_fault_library, fault): + """Rejected probes leave eager usable and preserve errors predating the probe.""" + import os + from pathlib import Path + import subprocess + import sys + import textwrap + + script = textwrap.dedent(""" + import ctypes + import os + import runpy + import sys + + import pytest + import torch + from freetoken.kernel import _cpu_moe + from freetoken.moe.cpu_executor import CpuMoeExecutor + + shim = ctypes.CDLL(os.environ["LD_PRELOAD"]) + stream_ptr = torch.cuda.current_stream().cuda_stream + fault = os.environ["FREETOKEN_TEST_HIP_PROBE_FAULT"] + if fault == "pending": + assert shim.probe_test_make_error() != 0 + pending = shim.probe_test_peek_error() + with pytest.raises(RuntimeError, match="pending error before probe"): + _cpu_moe.memops_probe(stream_ptr) + assert shim.probe_test_peek_error() == pending + assert shim.probe_test_begins() == 0 + sys.exit(0) + + helpers = runpy.run_path(sys.argv[1]) + torch.manual_seed(17) + cache = helpers["_make_cache"](2, 8, 256, 128, scale=0.03) + ex = CpuMoeExecutor( + cache, top_k=4, activation="silu", apply_router_weight_on_input=False, + num_threads=4, max_tokens=3, device=torch.device("cuda"), + ) + assert not ex._flag_sync and not ex._graph_capture_safe + assert shim.probe_test_injected() > 0 + assert shim.probe_test_begins() == shim.probe_test_ends() + assert shim.probe_test_peek_error() == 0 + if fault == "dependency": + assert shim.probe_test_removed_nodes() == 1 + assert shim.probe_test_last_end_rc() == 0 + else: + assert shim.probe_test_last_end_rc() != 0 + + for index in range(2): + torch.manual_seed(23 + index) + hidden = (torch.randn(3, 256) * 0.3).bfloat16().cuda() + ids = torch.randint(0, 8, (3, 4), dtype=torch.int32).cuda() + weights = torch.rand(3, 4).cuda() + actual = ex.decode(1, hidden, weights, ids).float().cpu() + expected = helpers["_reference"](cache, 1, hidden, ids.cpu(), weights.cpu()) + relative = (actual - expected).abs().max() / (expected.abs().max() + 1e-6) + assert relative < 2e-2, relative + + stream = torch.cuda.Stream() + with pytest.raises(RuntimeError, match="--cuda-graph-max-bs 0"): + with torch.cuda.graph(torch.cuda.CUDAGraph(), stream=stream): + ex.decode(1, hidden, weights, ids) + assert torch.isfinite(ex.decode(1, hidden, weights, ids).cpu()).all() + """) + env = dict(os.environ, LD_PRELOAD=str(hip_probe_fault_library), + FREETOKEN_TEST_HIP_PROBE_FAULT=fault, FREETOKEN_CPU_MOE_FLAG_SYNC="1") + result = subprocess.run( + [sys.executable, "-c", script, str(Path(__file__).resolve())], + env=env, capture_output=True, text=True, timeout=30, + ) + assert result.returncode == 0, result.stdout + result.stderr + + def test_cpu_moe_decode_cuda_graph_replay(): from freetoken.moe.cpu_executor import CpuMoeExecutor