perf(cpu-moe): extend expert dedup to ds_fp4 and mxfp4 - #49
Open
gdevenyi wants to merge 3 commits into
Open
Conversation
calvarado2004
added a commit
to calvarado2004/FreeToken
that referenced
this pull request
Aug 23, 2026
Add the exact DeepSeek-V4 TP4 CPU-MoE geometry to the community batch-sweep benchmark so the ds_fp4 dedup path is measured at the DSpark verification widths used in production, rather than extrapolated from batch size one. Report adaptive width and the selector's startup-profiled draft/verify costs without synchronizing the serving hot path. Document both the 28-30 tok/s high-survival result and the independently measured low-acceptance regression so deployments use representative A/B evidence instead of a universal threshold. Based-on: FlashML-org#41 Based-on: FlashML-org#49 Based-on: FlashML-org#81
gdevenyi
force-pushed
the
perf/cpu-moe-dedup-fp4
branch
from
August 23, 2026 16:19
38d3c9c to
a4d6ef2
Compare
Pass 1 split its work by (token, route, row block), so when two tokens in a decode
batch routed to the same expert, that expert's gate_up rows were streamed from DRAM
twice. The GEMV is DRAM-bound, so that is the whole cost.
Group the routes by expert (counting sort over the task's ids) and split pass 1 by
(unique expert, row block) instead, with the routes for that expert as the inner
loop. Each weight row is then read from memory once and reused from L1 across every
token routed to it -- both the gate row and the up row are 4 KiB at H=4096, so 8 KiB
stays resident across the inner loop. No kernel changes: the reuse comes from loop
order, so every format on the gemm1_dot path (bf16, nvfp4, fp8_block, q4_0) benefits
at once.
Measured on 2x Xeon Gold 6526Y, E=256, top_k=8, H=2048, I=768, uniform-random
routing (the pessimistic case -- real routing is skewed, so collisions are more
common):
bs routes unique reuse off on delta
4 32 30 1.07x 3.54ms 3.39ms +4.3%
8 64 56 1.14x 7.12ms 6.80ms +4.6%
16 128 104 1.23x 14.51ms 13.07ms +11.0%
32 256 156 1.64x 29.97ms 21.97ms +36.4%
64 512 213 2.40x 59.12ms 35.44ms +66.8%
Those track a simple traffic model to within 2%: pass 1 is about two thirds of the
bytes (gate_up is [2I, H] against down's [H, I]), so the expected speedup is
1 / (2/3 / reuse + 1/3) -- 1.35x at bs=32 and 1.64x at bs=64 against 1.36x and 1.67x
measured.
Inert below bs=2 and skipped when every route already has a distinct expert, so
single-stream decode keeps exactly the old work split. mxfp4 and ds_fp4 own their
pass-1 bodies and are untouched; they would follow the same shape.
Pass 2 is deliberately not deduped. Its work items are per-token and own their
output rows exclusively; deduping it would have several experts accumulating into
the same y row and needs a reduction, which is a separate change.
`FREETOKEN_CPU_MOE_DEDUP=0` restores the old split, and
`FREETOKEN_CPU_MOE_DEDUP_DEBUG=1` reports the reuse factor -- worth having, since a
first attempt at the toggle cached the env in a function-local static and silently
disabled both arms of the A/B, which read as "dedup does nothing" (+1.4%).
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GpGe2fQ5pDShGrnSuksfun
Pass 1 dedup left the down projection reading each expert once per route. Deduping
it the same way -- a work item per (expert, row block) -- would have several experts
summing into the same y row and need a cross-worker reduction.
Give one work item every token for its H-block instead. The rows are then owned
outright, the accumulation happens in a private fp32 buffer, and an expert's down
rows are still read once and reused across the tokens routed to it. The block is
HBLK_DD = 8 rather than 32 because the item count drops from tokens * n_hblk to
n_hblk, and H/8 keeps ~8 items per worker at H=2048 on a 32-core part.
On top of pass 1 (same rig: 2x Xeon Gold 6526Y, E=256, top_k=8, H=2048, I=768,
uniform-random routing):
bs reuse off pass1 pass1+pass2
8 1.14x 7.27ms +4.6% +11.8%
16 1.23x 14.63ms +11.0% +25.3%
32 1.64x 29.94ms +36.4% +78.8%
64 2.40x 59.18ms +66.8% +158.2%
2.58x at bs=64, slightly ahead of the 2.40x the traffic model predicts -- the
smaller H-block also helps locality.
Numerics: the fp32 accumulation now runs in expert order rather than route order, so
the last bits differ from the non-deduped path. That is the same latitude the kernel
already takes between ISA tiers, and the GPU-comparison test (bs 1/2/5/16) covers it.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GpGe2fQ5pDShGrnSuksfun
Both fp4 formats own their pass bodies, so neither got anything from the generic
dedup. Give them the same treatment: pass 1 splits by (unique expert, row block)
with the expert's routes inner, and pass 2 shares one deduped body -- one work item
owns an H-block for every token, accumulating in a private fp32 buffer.
ds_fp4 reuses its gate/up rows from L1 exactly as the generic path does. mxfp4's
mxgemv computes a whole tile per token, so its reuse is the tile staying resident
across the expert's routes -- 128 KiB at H=4096, L2 rather than L1, still not DRAM.
ds_fp4 (E=128 H=4096 I=2048 top_k=6) mxfp4 (E=64 H=2880 I=2880 top_k=4)
bs reuse off on delta bs reuse off on delta
8 1.20x 9.24ms 7.21ms +28.2% 8 1.39x 7.43ms 5.57ms +33.3%
16 1.35x 17.44ms 12.89ms +35.2% 16 1.45x 13.68ms 10.72ms +27.6%
32 1.94x 33.86ms 19.62ms +72.5% 32 2.33x 26.57ms 14.70ms +80.7%
The deduped pass-2 block size has to be format-aware, and getting it wrong is
expensive. For the row-major formats the block is just "which output rows", so 8
rows is free and keeps the worker pool fed now that the item count has dropped from
tokens * n_hblk to n_hblk. mxfp4's bank is transposed, so the same number becomes
mxgemv's `ncol` -- the dimension it vectorizes over -- and below 16 every call falls
into mxgemv's scalar tail. Measured at HBLK_DD = 8, mxfp4 dedup ran 2.7x SLOWER than
not deduping (-63% at bs=8, -62% at bs=32). It keeps the full 32-row block.
Covered by the existing GPU-comparison tests, which run these formats at batch sizes
where dedup engages: test_cpu_decode_mxfp4_matches_gpu_splitk at bs 1/4/8 and
test_cpu_decode_dsfp4_matches_gpu at bs 1/3/8.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GpGe2fQ5pDShGrnSuksfun
gdevenyi
force-pushed
the
perf/cpu-moe-dedup-fp4
branch
from
August 23, 2026 16:24
a4d6ef2 to
4a6d86b
Compare
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
1 participant
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.
What this adds
#46 and #47 deduplicate expert reads for bf16 experts. The
ds_fp4andmxfp4paths have their own inner loops and did not get the same treatment.This PR extends the deduplicated walk to both. Each format keeps its own block-scale handling; only the route walk changes.
Stacked on #47 and #46
This PR contains both. All three edit the same route-walking code in
cpu_moe_ext.cpp, and this one reuses the dedup list they build. Review #46, then #47, then this.Measurements
Measured on
mainwith only these three PRs, using the batch sweep from #81. The bf16 path gives:The bf16 rows match #47 within run-to-run spread, which is what should happen: this PR changes the fp4 paths, not bf16. Its value is that
ds_fp4andmxfp4deployments get the same behaviour instead of falling back to the per-route walk.Testing
tests/moe/test_cpu_moe.pyonmainwith this PR: 20 passed.This PR adds no test of its own. The fp4 dedup paths are covered only by the existing correctness tests for those formats, which do not exercise a batch large enough to force reuse. A test that pins the fp4 dedup output against the per-route output at batch 32 would close that gap, and I am happy to add one if you want it before merge.