Repository navigation
feat(cpu-moe): serve block-fp8 experts on the cpu and hybrid backends - #662
Merged
Merged
Conversation
Block-fp8 checkpoints (Qwen3.5/3.6-FP8, DeepSeek-V3-style, GLM) had no CPU expert
path at all:
NotImplementedError: --moe-backend cpu/hybrid ... supports
['bf16', 'ds_fp4', 'mxfp4_triton', 'nvfp4', 'q4_0'] formats,
but this checkpoint's experts are 'fp8_block'
so every expert miss had to stream over PCIe at ~25 GB/s while the CPU sat idle.
The bank is fp8-e4m3 row-major [rows, K] with one bf16 scale per 128x128 block --
K contiguous per output row, unlike mxfp4's transposed bank, so it maps straight
onto a bf16 dot product. e4m3 -> bf16 is *exact* (3 mantissa bits into 7, and
bf16's exponent range covers e4m3's whole 2^-9..448 span), so the AVX-512 path
widens weights in-register and feeds `_mm512_dpbf16_ps` against the bf16
activations: identical products to the reference's fp32 math, differing only in
summation order -- the same latitude the bf16 dot already takes.
The normal range rebases the exponent with a shift and a constant +120; exp == 0
is subnormal (m * 2^-9) and does not obey that, so those lanes blend in from a
small table. AVX2 does the same in fp32; scalar goes through the existing 256-entry
e4m3 LUT.
`ft bench bw` gains an fp8_block CPU measurement (it could already build the banks).
Measured on main with this commit alone, `--dtype fp8_block --reps 5`, 2x Xeon Gold
6526Y:
CPU MoE PCIe ratio backend picked
main not measurable 25.1 -- offload
main + this 53.7 GB/s 25.1 2.14x hybrid
Before this commit the format reports no CPU number, so the backend picker has
nothing to weigh and selects offload. The ratio is what decides the backend, so
adding the kernel is what lets a block-fp8 deployment use the CPU at all.
End to end on Qwen/Qwen3.6-35B-A3B-FP8, median of 3 generations at matched GPU
expert residency:
residency offload hybrid delta
10% 43.94 50.94 +15.9%
25% 82.65 55.12 -33%
The win is real but conditional: hybrid pays off once the miss rate is high enough
that PCIe is the bottleneck, and costs throughput when enough experts are resident
that the CPU becomes the limiter. See the PR discussion -- the ratio makes
`--moe-backend auto` select hybrid for every block-fp8 checkpoint, which is right
for the oversubscribed case this project targets and wrong for a model that nearly
fits, and the bench models bandwidth but not miss rate.
Correctness is checked against the Triton block-fp8 decode kernel on identical
banks and routing (bs 1/2/5), plus an exactness check on the e4m3 -> bf16 widening
over all 256 codes. A weight-format ABI probe (`max_weight_format_id`) makes a
stale prebuilt `.so` fail loudly instead of falling through to the wrong dequant
branch and returning silently wrong numbers.
Assisted-by: Claude
The scale banks' trailing dimension is padded so per-expert rows land on a 16-byte boundary for the fused copy (kernel/aot_models.py: fp8_block_scale_pad), so the row stride is not ceil(K/128). Recomputing it from the unpadded block count reads scales at the wrong offset. Qwen3.8 is the case in the tree: models/qwen3_5_moe/weight.py records down_scale as 20x5 bf16 = 200 B, and 200 % 16 != 0 pads cols 5 -> 6. The benchbw workload (H=2048, I=512) and DSV4-Flash (H=7168, I=2048) both happen to give rows*cols divisible by 8, so neither trips it. Read the stride off the tensor and pass it through to the kernel, and let the shape check assert the unpadded extent is present rather than an exact shape. The fixture now allocates padded banks, so the existing GPU cross-check covers this too; the GPU kernel was already stride-aware.
The existing GPU-vs-CPU check compares two fp32 accumulations at a 2e-2 tolerance, which cannot see accumulator precision loss. Add a float64 reference over mixed-sign, full-range e4m3 weights and zero-mean activations, and force each ISA tier explicitly so the sweep is real. Also add an avx512f tier between avx512bf16 and avx2, so CPUs with AVX-512 but no BF16 extension (Skylake-SP through Ice Lake-SP) stop falling back to 8-wide fp32. Measured on Xeon Platinum 8368 (Ice Lake-SP, no avx512_bf16), three repetitions each: +19% single-thread cache-resident (3.31 -> 3.94 GB/s, disjoint intervals), but only +4% under `ft bench bw --dtype fp8` at 75 threads (60.4 -> 62.7 GB/s, overlapping intervals). At that thread count the kernel is memory-bound, so the project's own benchmark cannot resolve the tier; the single-thread gain is the honest figure.
The new tier shipped without test coverage: deleting the avx512f branch from select_fp8dot left every test green, because isa_name() reports select_dot() (the bf16 chain) rather than the fp8 dispatch. Injecting a ragged-tail bug into dot_fp8_block_avx2 was invisible for the same reason -- no test ran numerics under a forced tier. Force each available fp8 tier through the float64 parity GEMV, expose the selected fp8 tier so the assertion can see it, and make the padded bank test run a GEMV instead of only constructing an executor. Also move the module-level CUDA guard onto the tests that actually need a GPU. A CPU MoE kernel's most relevant hosts are the ones without a usable GPU, and all ten tests were skipping there.
…ring overclaiming The comparison test only varied batch size, so all three cases ran whatever tier this machine auto-selects and the scalar / AVX2 / AVX-512-BF16 fp8_block kernels were never compared against the GPU at all. Parametrise FREETOKEN_CPU_MOE_ISA over the four tiers. Setting the variable alone would not be enough: pick_isa() caps *downward*, so asking for a tier the CPU or build lacks silently returns a lower one and the test would pass four times over the same kernel. The executor latches its choice into isa_name() at construction, so read it back and skip when it does not match what was asked for. On a machine with all four, that is 4 x 3 cases against a 8-expert, 256-hidden problem. test_subnormal_and_sign_decode_exactly said it validated that "every e4m3 byte must widen to the value the reference LUT gives", but it never calls the CPU kernel -- it checks that bf16 round-trips e4m3 losslessly and that subnormals are finite and small. That is the premise the widening rests on, which is worth pinning, but it is not the widening. Renamed to test_bf16_carries_every_e4m3_byte and the docstring now says what the body asserts. Not run here: this file is skipif(not torch.cuda.is_available()), and this machine has no GPU and no torch. Compiles clean and ruff reports the same two pre-existing E741 findings before and after. Assisted-by: Claude
KarrAcaRn
pushed a commit
to KarrAcaRn/FreeToken-ByAI
that referenced
this pull request
Oct 12, 2026
Brings upstream FlashML-org#662 (block-fp8 CPU MoE) and FlashML-org#527 (sampling penalties). - FlashML-org#662 is the merged form of FlashML-org#399, already on next: take upstream's kernel, executor, benchbw and test (fp8_block_* names, profile version check) and keep next's ROCm memop path in cpu_moe_ext.cpp. test_hybrid_fetch's geometry test now writes a versioned profile. - FlashML-org#527: keep next's implementation (5a4987c): counts rebuilt from token_pool each step, so DFlash multi-token accepts and overlap scheduling stay right, and out-of-range penalties stay a 400 with param instead of a 422. - docs/cli.md: next's flags, with fp8_block in the --moe-cpu-layers formats. Assisted-by: Claude Opus 5.5
KarrAcaRn
pushed a commit
to KarrAcaRn/FreeToken-ByAI
that referenced
this pull request
Oct 12, 2026
FlashML-org#665 adopted with a negative-slot test, FlashML-org#378 deferred again (ROCm, now ready); FlashML-org#399 landed upstream via FlashML-org#662. Assisted-by: Claude Opus 5.5
3 tasks done
trcwebdesign
pushed a commit
to trcwebdesign/FreeToken
that referenced
this pull request
Oct 14, 2026
…FlashML-org#662) * feat(cpu-moe): block-FP8 expert GEMV (W8A16) for the cpu/hybrid backends Block-fp8 checkpoints (Qwen3.5/3.6-FP8, DeepSeek-V3-style, GLM) had no CPU expert path at all: NotImplementedError: --moe-backend cpu/hybrid ... supports ['bf16', 'ds_fp4', 'mxfp4_triton', 'nvfp4', 'q4_0'] formats, but this checkpoint's experts are 'fp8_block' so every expert miss had to stream over PCIe at ~25 GB/s while the CPU sat idle. The bank is fp8-e4m3 row-major [rows, K] with one bf16 scale per 128x128 block -- K contiguous per output row, unlike mxfp4's transposed bank, so it maps straight onto a bf16 dot product. e4m3 -> bf16 is *exact* (3 mantissa bits into 7, and bf16's exponent range covers e4m3's whole 2^-9..448 span), so the AVX-512 path widens weights in-register and feeds `_mm512_dpbf16_ps` against the bf16 activations: identical products to the reference's fp32 math, differing only in summation order -- the same latitude the bf16 dot already takes. The normal range rebases the exponent with a shift and a constant +120; exp == 0 is subnormal (m * 2^-9) and does not obey that, so those lanes blend in from a small table. AVX2 does the same in fp32; scalar goes through the existing 256-entry e4m3 LUT. `ft bench bw` gains an fp8_block CPU measurement (it could already build the banks). Measured on main with this commit alone, `--dtype fp8_block --reps 5`, 2x Xeon Gold 6526Y: CPU MoE PCIe ratio backend picked main not measurable 25.1 -- offload main + this 53.7 GB/s 25.1 2.14x hybrid Before this commit the format reports no CPU number, so the backend picker has nothing to weigh and selects offload. The ratio is what decides the backend, so adding the kernel is what lets a block-fp8 deployment use the CPU at all. End to end on Qwen/Qwen3.6-35B-A3B-FP8, median of 3 generations at matched GPU expert residency: residency offload hybrid delta 10% 43.94 50.94 +15.9% 25% 82.65 55.12 -33% The win is real but conditional: hybrid pays off once the miss rate is high enough that PCIe is the bottleneck, and costs throughput when enough experts are resident that the CPU becomes the limiter. See the PR discussion -- the ratio makes `--moe-backend auto` select hybrid for every block-fp8 checkpoint, which is right for the oversubscribed case this project targets and wrong for a model that nearly fits, and the bench models bandwidth but not miss rate. Correctness is checked against the Triton block-fp8 decode kernel on identical banks and routing (bs 1/2/5), plus an exactness check on the e4m3 -> bf16 widening over all 256 codes. A weight-format ABI probe (`max_weight_format_id`) makes a stale prebuilt `.so` fail loudly instead of falling through to the wrong dequant branch and returning silently wrong numbers. Assisted-by: Claude * fix(cpu-moe): honor padded fp8 scale strides The scale banks' trailing dimension is padded so per-expert rows land on a 16-byte boundary for the fused copy (kernel/aot_models.py: fp8_block_scale_pad), so the row stride is not ceil(K/128). Recomputing it from the unpadded block count reads scales at the wrong offset. Qwen3.8 is the case in the tree: models/qwen3_5_moe/weight.py records down_scale as 20x5 bf16 = 200 B, and 200 % 16 != 0 pads cols 5 -> 6. The benchbw workload (H=2048, I=512) and DSV4-Flash (H=7168, I=2048) both happen to give rows*cols divisible by 8, so neither trips it. Read the stride off the tensor and pass it through to the kernel, and let the shape check assert the unpadded extent is present rather than an exact shape. The fixture now allocates padded banks, so the existing GPU cross-check covers this too; the GPU kernel was already stride-aware. * test(cpu-moe): add float64 parity reference and an avx512f fp8 tier The existing GPU-vs-CPU check compares two fp32 accumulations at a 2e-2 tolerance, which cannot see accumulator precision loss. Add a float64 reference over mixed-sign, full-range e4m3 weights and zero-mean activations, and force each ISA tier explicitly so the sweep is real. Also add an avx512f tier between avx512bf16 and avx2, so CPUs with AVX-512 but no BF16 extension (Skylake-SP through Ice Lake-SP) stop falling back to 8-wide fp32. Measured on Xeon Platinum 8368 (Ice Lake-SP, no avx512_bf16), three repetitions each: +19% single-thread cache-resident (3.31 -> 3.94 GB/s, disjoint intervals), but only +4% under `ft bench bw --dtype fp8` at 75 threads (60.4 -> 62.7 GB/s, overlapping intervals). At that thread count the kernel is memory-bound, so the project's own benchmark cannot resolve the tier; the single-thread gain is the honest figure. * test(cpu-moe): cover fp8 dispatch and cpu-only regressions The new tier shipped without test coverage: deleting the avx512f branch from select_fp8dot left every test green, because isa_name() reports select_dot() (the bf16 chain) rather than the fp8 dispatch. Injecting a ragged-tail bug into dot_fp8_block_avx2 was invisible for the same reason -- no test ran numerics under a forced tier. Force each available fp8 tier through the float64 parity GEMV, expose the selected fp8 tier so the assertion can see it, and make the padded bank test run a GEMV instead of only constructing an executor. Also move the module-level CUDA guard onto the tests that actually need a GPU. A CPU MoE kernel's most relevant hosts are the ones without a usable GPU, and all ten tests were skipping there. * review: force each ISA tier in the fp8_block test, and stop the docstring overclaiming The comparison test only varied batch size, so all three cases ran whatever tier this machine auto-selects and the scalar / AVX2 / AVX-512-BF16 fp8_block kernels were never compared against the GPU at all. Parametrise FREETOKEN_CPU_MOE_ISA over the four tiers. Setting the variable alone would not be enough: pick_isa() caps *downward*, so asking for a tier the CPU or build lacks silently returns a lower one and the test would pass four times over the same kernel. The executor latches its choice into isa_name() at construction, so read it back and skip when it does not match what was asked for. On a machine with all four, that is 4 x 3 cases against a 8-expert, 256-hidden problem. test_subnormal_and_sign_decode_exactly said it validated that "every e4m3 byte must widen to the value the reference LUT gives", but it never calls the CPU kernel -- it checks that bf16 round-trips e4m3 losslessly and that subnormals are finite and small. That is the premise the widening rests on, which is worth pinning, but it is not the widening. Renamed to test_bf16_carries_every_e4m3_byte and the docstring now says what the body asserts. Not run here: this file is skipif(not torch.cuda.is_available()), and this machine has no GPU and no torch. Compiles clean and ruff reports the same two pre-existing E741 findings before and after. Assisted-by: Claude * fix(moe): let fp8_block experts decode on the cpu executor * fix(cpu-moe): build the bf16 fp8 tier only when the compiler has avx512bf16 * fix(bench): version the bw profile and ignore older ones * perf(cpu-moe): read fp8 weights as fp16 bits in the cpu dot kernels * refactor(cpu-moe): write the fp8_block bank checks like the other formats * refactor(cpu-moe): name the fp8 cpu symbols fp8_block * docs(server): list fp8_block in the --moe-cpu-layers help --------- Co-authored-by: Gabriel A. Devenyi <gdevenyi@gmail.com> Co-authored-by: ChenyuHe <hechenyu@zju.edu.cn>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Supersedes #399, which builds on #36 (closed). The commits of @gdevenyi (#36, including 19e9382) and @ChenyuHeee (#399) are kept to preserve their credit. The commits after them address the review on #399 and rewrite the fp8_block dot kernels.
Block-FP8 checkpoints (e4m3 experts, one bf16 scale per 128x128 block) can now run
--moe-strategy cpuand--moe-strategy hybrid. Onmain,--moe-strategy hybridfails at kernel selection:KernelSelectionError: no usable kernel in table; triton: has no CPU executor format; decode must run on the GPU.Related: #534 (the same
has no CPU executor formatrejection undermoe_strategy=offload).Call path:
OffloadMoELayer._decode_routed->CpuMoeExecutor.decode(cpu layers) or_decode_hybrid->CpuMoeExecutor.decode_submit/decode_sync->run_task_body->do_pass1/do_pass2->gemm1_dot/gemm2_dot->dot_fp8_block_{avx512f,avx2,scalar}. Decode only; prefill stays on the GPU. W8A16, the same math as the GPU decode kernel_decode_fp8_moe_kernel.Kernels
(b & 0x80) << 8 | (b & 0x7F) << 7is its e4m3 value times 2^-8, subnormals included, sovcvtph2psdecodes it and the block scale carries the 256 (same bit trick as the Marlin fp8 dequant in vLLM). Each 16-bit lane holds two bytes;srai 1 & 0xBF80builds the odd one, a left shift by 8 first builds the even one, so the activations are read as fp32 even/odd halves. The executor already produces those halves for nvfp4 and ds_fp4 (needs_di); fp8_block now uses them too. Four accumulators, block scale applied with a vector FMA, one horizontal sum per row.float8_e4m3fn) decode to +-480 in every kernel, as before.Other changes on top of #399
TritonFp8BlockMoEKerneldeclarescpu_format = "fp8_block"and accepts a CPU decode target.ft bench bwprofiles carry version 5. Older profiles are ignored with a warning to rerunft bench bw: they recordedfp8_block -> offloadwith no CPU number. Until the bench is rerun, auto stays on offload for every format and hybrid uses a fetch cap of 1.cpu_executor.pyfollow the other formats.Bandwidth
H100 80GB HBM3 + Xeon Platinum 8480+ (13 physical cores visible to the process), torch 2.14.1, gcc 11.4. Old = the kernels from #399 (7347a0e in this branch), new = this PR.
ft bench bw --dtype fp8 --model qwen3.6-moe --formats fp8_block --isa all, median of 3 runs per side, GB/s:--dtype fp8old -> newqwen3.6-moeold -> newMore shapes,
measure_cpu_moe_bwfromfreetoken/moe/benchbw.py(bs 1, top-8, E=256), median of 2 runs for fp8_block and 1 run for bf16. GB/s, new fp8_block (old fp8_block):An fp8_block expert is half the bytes of a bf16 one, so at these rates an expert decodes 1.4x to 2.5x faster in fp8_block than in bf16 at the same tier. With the old kernels, avx2 was slower per expert than bf16 at every shape and avx512f was within about 10% of bf16.
End to end
Run before the kernel rewrite (with the kernels from #399), Qwen/Qwen3.6-35B-A3B-FP8,
ft serve --model-path <ckpt> --moe-strategy <strategy> --moe-cache-rate 0.25. Greedy: 4 short prompts, temperature 0, thinking off, up to 256 tokens. AIME25 problem 1: the checkpoint's sampling defaults, up to 16384 tokens.Tests
tests/moe tests/models/test_quant_config.pyon the H100 box at the head of this PR: 180 passed, 39 skipped, 1 failed. The failure istest_scheme_for_agrees_with_the_stored_tensors[apodex/Apodex-1.1-mini-NVFP4[index-only]], a local HF cache entry with an unsupportedmodelopt_mixedconfig.tests/moe/test_cpu_moe_fp8_block.pywithout CUDA: 14 passed, 9 skipped (the GPU parity cases). The float64 reference case runs every forced tier on full-range codes and on subnormal-only codes with ragged K.cpu_moe_ext.cppbuilds with gcc 11.4 without warnings and passesclang++-19 -fsyntax-only -Wall -Wextra.