perf(cpu-moe): dedup experts in pass 2 as well - #47
Open
gdevenyi wants to merge 2 commits into
Open
Conversation
This was referenced Aug 22, 2026
gdevenyi
force-pushed
the
perf/cpu-moe-dedup-pass2
branch
from
August 23, 2026 16:19
5483c35 to
9883682
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
gdevenyi
force-pushed
the
perf/cpu-moe-dedup-pass2
branch
from
August 23, 2026 16:24
9883682 to
a60a654
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 deduplicates the expert reads in pass 1. Pass 2 still reads
downonce per route.This PR reuses the same deduplicated route list for pass 2. Pass 2 accumulates into a per-token output rather than reading a shared operand, so it needs a scatter step that #46 did not.
Stacked on #46
This PR contains #46's commit. The two touch the same code in
cpu_moe_ext.cppand pass 2 reuses the route list that #46 builds, so they cannot be separated without duplicating that work. Review #46 first.Measurements
Measured on
mainwith this PR alone, using the batch sweep from #81 as the measuring tool. Same machine and workload as #46, so the rows compare directly.At batch 64 the cost stops rising with the batch: 9.91 ms against 9.81 ms at batch 32, while the route count doubles from 256 to 512. Once both passes read each distinct expert once, the step costs what the 126 distinct experts cost, and adding more tokens that select the same experts is nearly free.
That is the point of the pair. #46 alone reaches 2.07x at batch 64 because it still pays full price for one third of the bytes. Removing that third takes the same workload to 4.00x.
Testing
tests/moe/test_cpu_moe.pyonmainwith this PR: 20 passed.