perf(cpu-moe): dedup experts across the batch in pass 1 - #46
Conversation
|
Re-reading the paper (arXiv:2608.16157) I want to head off a likely objection, because §4.1 already contains the word "deduplicates":
That is a different dedup, on the other side of the boundary:
The CPU executor never sees the deduplicated set. §4.1 says the kernel "rewrites logical routed IDs into physical slot IDs or a special CPU-assignment flag", so what reaches the workers is still per-route: So the two compose rather than overlap — §4.1 stops the expert crossing PCIe twice, this stops it crossing the memory bus twice once it is already in host RAM. This also lines up with the paper's own characterisation of the CPU path in the same section:
Which is exactly why the win here tracks the traffic model to within 2% and why it comes from loop order rather than any kernel change. |
b8a30ee to
126fe5c
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
126fe5c to
49d8a0f
Compare
1 participant
What this adds
The CPU MoE decode kernel walks the routed experts one route at a time. When several tokens in a batch route to the same expert, it reads that expert's weights once per token.
This PR builds a deduplicated route list before pass 1 and reads each distinct expert once, accumulating into every token that selected it. The weights are the large operand, so the saving tracks how often experts repeat.
Measurements
Measured on
mainwith this PR alone, using the batch sweep from #81 as the measuring tool. Two Xeon Gold 6526Y, bf16 experts of 9.00 MB, 256 experts, top-6.Both arms replay the same routing draw, so the comparison uses identical work.
At batch 1 there is nothing to deduplicate. The 1.11x there is the deduplicated path being slightly cheaper to walk, not reuse.
Why the numbers land where they do
Pass 1 reads
gate_up, shaped[2I, H]. Pass 2 readsdown, shaped[H, I]. Pass 1 is therefore about two thirds of the bytes. This PR deduplicates pass 1 only, so reuse alone predicts1 / (2/3 / reuse + 1/3):Measured is above predicted at both sizes, and the batch-1 row says why. At batch 1 nothing repeats, reuse is 1.00x, and the formula predicts no gain at all, yet the step still runs 1.11x faster. So the deduplicated walk is itself cheaper than the per-route walk, independent of reuse.
Folding that constant in gives 1.78x and 2.23x, which now sit slightly above the measured 1.74x and 2.07x. The measurements land between the two estimates. I read that as the reuse term carrying the effect, with a small fixed saving on top, rather than as either model being exact.
#47 deduplicates pass 2 as well and takes the remaining third.
Cost and safety
The dedup list is built once per layer per step.
FREETOKEN_CPU_MOE_DEDUP=0restores the old path for A/B testing.Testing
tests/moe/test_cpu_moe.pyonmainwith this PR: 20 passed.