perf(moe): keep the CPU MoE pool and its expert banks on one NUMA node

[gdevenyi]

Problem

cpu_moe_ext.cpp:1094 states the assumption outright — "NUMA: a single node is assumed" — but nothing enforced it. resolve_threads_and_affinity(0) returned every physical core on the machine, so on a 2-socket box half the workers read expert bytes over the interconnect, and the sense-reversing barrier plus the p1_next/p2_next work counters ping-pong their cache lines across it.

On a 2× Xeon Gold 6526Y the pool reached 26% of the STREAM ceiling spanning both sockets, against 91% confined to one.

Results

ft bench bw, median of 5, no numactl — this is the default-path delta a user sees:

dtype	before	after	delta
bf16	67.5	96.2	+42.5%
ds_fp4	66.9	86.9	+29.9%
mxfp4	48.5	59.8	+23.3%
nvfp4	64.4	75.7	+17.5%

Hardware: 16c/socket, 4 of 8 DDR5-5200 channels populated per socket, 2× RTX 6000 Ada both on node0, Ubuntu 24.04, torch 2.11.0+cu130.

Two halves, and neither works alone

Confining the workers while the banks stay unplaced made things worse, not better — throughput went bimodal, 56–123 GB/s run to run, against a steady 69 before. The banks were landing wherever the loader threads happened to run. Both halves together are stable: bf16 spans 94.2–99.0 across five runs.

• Threads — moe_pool_numa_node() picks the GPU's node, keeping the expert GEMV and the offload gather's DMA local. Confining would otherwise halve the pool, so the default there fills the node's SMT siblings too: the thread count is unchanged and only the placement moves. That matters for the compute-bound formats — at one-thread-per-core on half the cores mxfp4 lost 23%, and SMT brings it back to +23%.
• Banks — HostBank mmaps get mbind(MPOL_PREFERRED) before the fill faults them in. Pin-after-fill means nothing is resident yet, which is the only moment placement is free. A preference, not a reservation: the banks are 130+ GiB and MPOL_BIND would OOM rather than spill once the node fills — a far worse failure than a remote read.

mbind goes through syscall() rather than libnuma (glibc does not wrap it), so no new runtime dependency. Verified at page level: prefer:1 … N1=16384.

benchbw places its synthetic banks with set_mempolicy instead — it allocates via cudaHostAlloc, which faults and pins up front, and pinned pages can never be migrated.

Not breaking single-socket

FreeToken's usual target is a single-socket desktop, where this must be a no-op. With fewer than two nodes the resolver returns exactly what it always did and every placement call short-circuits. Same on non-Linux, on an unknown syscall ABI, and under FREETOKEN_CPU_MOE_NUMA=off.

test_single_node_is_not_confined, test_single_node_pool_is_unchanged, test_no_node_argument_matches_old_behaviour and test_unreadable_topology_is_not_confined are the regression guards.

One subtlety worth flagging for review: the torch intra-op clamp at cpu_executor.py keeps reading physical_core_cpus(pool_node) rather than the machine-wide count, so it lands exactly where it did before. Left as-is deliberately — retuning it is a separate question.

Escape hatch

FREETOKEN_CPU_MOE_NUMA takes auto (default), off, or a node id. No boolean spellings on purpose: 0 and 1 are node ids, so on/off would make the common case ambiguous.

Tests

20 new tests in tests/moe/test_cpu_pool_numa.py, including one that mbinds a real mapping and asserts the policy takes (skipped on single-node machines).

Full suite on the target: 743 passed, 6 skipped, 1 failed — the failure is tests/moe/test_cpu_moe_q4_0.py::test_cpu_decode_q4_0_matches_ggml_mmvq, which also fails on main at the same commit and is unrelated.

ruff check introduces no new violations on any modified file (counts identical to HEAD).

Follow-ups (not in this PR)

• mxfp4 still runs the W4A16 permutexvar+fp32 path while nvfp4 has a W4A8 VNNI one, using the same e2m1 codes — it is the only compute-bound format and the reason it needs SMT to break even.
• The DRAM ceiling on this box is halved by 4-of-8 channels populated; that is hardware, not code.

🤖 Generated with Claude Code

https://claude.ai/code/session_01GpGe2fQ5pDShGrnSuksfun

[gdevenyi]

Correction to the first follow-up bullet above, before it misleads anyone.

I wrote that mxfp4 could take nvfp4's W4A8 VNNI path since both use the same e2m1 codes. That is wrong — the codes match but the bank layouts are orthogonal:

format	bank	contiguous dim
nvfp4	row-major [rows, K]	K
mxfp4	blocks_t [K//2, N] (transposed split-K)	N (output columns)

VPDPBUSD reduces four adjacent int8 pairs into one int32. For nvfp4 those are adjacent K, which is the dot product. For mxfp4 adjacent bytes are four different output columns, so the same instruction would sum across columns. _mm512_dpbf16_ps fails the same way.

Making VNNI usable there needs blocks_t transposed to K-contiguous, and that bank is shared with the Triton GPU decode kernel and the transposed grouped prefill path — a much bigger change than a kernel port.

The observation that stands is that mxfp4 is the only compute-bound format (59.8 GB/s vs bf16's 96.2 on more bytes per expert), and that its decode chain consumes only 16 of 64 byte lanes per load. Options that stay inside the current layout — 64-byte _mm512_shuffle_epi8 decode, or a bf16 LUT via _mm512_permutexvar_epi16 — are unmeasured and not a quick win.

None of this affects the change in this PR; mxfp4 is +23.3% here purely from placement.

[gdevenyi]

Moving this to draft. Do not merge — it is a net regression in real serving, and the numbers in the description are misleading.

I finally measured end to end instead of trusting ft bench bw, on DeepSeek-V4-Flash-0731 (ds_fp4 experts, 143 GiB of banks, --moe-backend hybrid), alternating configs across two full cycles, 51 decode-throughput samples per run:

config	cycle 1	cycle 2
this PR (FREETOKEN_CPU_MOE_NUMA=auto)	17.93 tok/s	18.07 tok/s
today's behaviour (=off)	19.21 tok/s	19.39 tok/s

~6.7% slower, reproducible, order-controlled, page cache dropped beforehand and both NUMA nodes balanced at ~251 GB free at the start.

The microbenchmark in the description says +29.9% for this same format on this same machine. Both numbers are real; the microbenchmark simply measures the wrong thing.

Why

ft bench bw times the CPU MoE GEMV essentially in isolation. Real hybrid decode runs it concurrently with the PCIe expert gather (34.1% of each step's misses here) plus KV/attention traffic and the CUDA driver threads. This box has only 4 of 8 DDR5 channels populated per socket, so one node is ~166 GB/s and the pair is ~333 GB/s. Confining every consumer to node0 makes locality better and total available bandwidth half what it was. For the isolated GEMV that trade wins; for the real concurrent workload it loses.

Confinement also puts 31 workers on node0's 16 physical cores while the flag-sync coordinator, the API server, and the driver threads all now contend for the same node, where before they had a whole idle socket.

What I think is salvageable

The machinery (topology detection, mbind(MPOL_PREFERRED) for pin-after-fill banks, the escape hatch, the single-node no-op guarantees) is sound and tested. The wrong part is defaulting it on.

The untested hypothesis worth checking before anyone revives this: under --moe-backend cpu, where there is no concurrent PCIe gather competing for DRAM, confinement should win — that is the configuration the microbenchmark actually resembles. If that holds, the right shape is default-off with FREETOKEN_CPU_MOE_NUMA=<node> as an opt-in, documented as helping CPU-dominated backends and hurting hybrid on bandwidth-starved multi-socket boxes.

I have not tested that yet, so I am not re-proposing it — just recording where the evidence points. Apologies for the noisy description; the +42.5%/+29.9% figures should be read as "isolated CPU GEMV only", not as serving throughput.

[gdevenyi]

Tested the hypothesis from my last comment. It is refuted, and worse than I expected — so I am closing this.

I guessed that confinement would win under --moe-backend cpu, where no concurrent PCIe gather competes for DRAM. Same alternating two-cycle harness, DeepSeek-V4-Flash-0731, 19 decode samples per run:

--moe-backend cpu	cycle 1	cycle 2
confined (FREETOKEN_CPU_MOE_NUMA=auto)	8.01 tok/s	7.91 tok/s
today's behaviour (=off)	11.21 tok/s	11.02 tok/s

28% slower — four times the penalty seen with hybrid (6.7%). Removing the competing PCIe traffic made confinement worse, not better, which is the exact opposite of the prediction.

What is actually going on

Pure-CPU expert decode is the most DRAM-hungry configuration there is: every expert byte of every layer is streamed from host memory, read once, never reused. What matters for a read-once streaming workload is aggregate memory-controller bandwidth, not locality. Spanning both sockets gets ~333 GB/s across two memory controllers; confining to one node gets ~166 GB/s. A remote read over UPI is cheaper than losing half your bandwidth.

That also explains why ft bench bw was so badly misleading: its synthetic banks are ~2 GiB and one layer deep, so once confined and mbind-ed they sit entirely in the local node and are re-read hot. The real path streams 143 GiB.

So cpu_moe_ext.cpp's "NUMA: a single node is assumed" is not a latent bug to fix. Spreading the pool over every core is the right default, and this PR was solving a problem that does not exist.

Closing. Sorry for the noise — the microbenchmark numbers in the description were real but measured something that does not predict serving throughput.

What might actually be worth doing

Not confinement — interleaving. Bank pages currently land wherever the 8 unpinned loader threads happen to run, so the split between nodes is a lottery. MPOL_INTERLEAVE would guarantee an even 50/50 and keep both memory controllers evenly loaded. That is the opposite of what this PR does, it is untested, and I am not proposing it without end-to-end numbers.