refactor(cpu-moe): split the executor into weight formats and gpu sync backends

[jason-fxz]

Summary

cpu_moe_ext.cpp held the CPU MoE executor, the kernels of all six weight formats and the CUDA-only GPU/CPU handshake in one 2.4k-line file. This PR splits it so that a weight format, an ISA tier or a GPU runtime can change without touching the executor. It also moves the GPU side of the handshake into two kernels on CUDA, and the coordinator skips tasks without routed experts (see Kernel IO).

Layout of kernel/csrc/cpu_moe/:

File	Role
cpu_moe_ext.cpp	Executor (worker pool, pass 1 / pass 2 work split, coordinator, host-function fallback) and the pybind module
format.h	MoeFormat interface, MoeParams (shared shape and activation fields), FormatArgs (banks and options by name)
kernels/<fmt>.cpp	One file per format: its ISA kernels, kernel selection and MoeFormat class (bf16, nvfp4, mxfp4_triton, ds_fp4, q4_0, fp8_block)
common.h	ISA tier selection, CPU_MOE_TARGET, numeric helpers, 64-byte aligned scratch
gpu_sync.h, gpu_sync_cuda.cpp, gpu_sync_hip.cpp	GPU/CPU handshake; setup.py builds the backend of the installed torch
gpu_sync_io.h, gpu_sync_io.cu	Kernel IO for the handshake; one source, built by nvcc on CUDA and by hipcc on ROCm through freetoken/hip_compat.h

Format interface

• A format implements gate_up(layer, e, tok, i0, i1, x, w_in, g) for a block of intermediate rows and down(layer, h0, h1, routes, n, out) for a block of output rows over all of a token's routes at once. Visiting the routes' rows together keeps the down GEMV at DRAM bandwidth.
• prep_input, prep_inter and rounds_routes are optional hooks with defaults.
• A format reads the banks and integer options it needs by name in its constructor. MoeParams holds only the fields every format shares.
• The formats are listed once, by name, in kFormats; _cpu_moe.weight_formats() exports the names. This replaces the integer format ids kept in sync between C++ and Python.

Handshake

• The GPU writes one doorbell word MoeTask{layer:32, num_tokens:32} into a single ready word and waits on a single done word. Inputs and outputs go through one shared pinned IO block of max_tokens rows. This replaces the per-(layer, batch size) flag slots, pinned IO buffers and task table. It holds because one task is in flight: decode runs on one stream and every layer waits before the next one submits.
• GpuSync states two contracts a backend keeps in eager mode and in graph replay: submit (GPU to CPU) and wait (CPU to GPU).
• CUDA backend: the existing cuStreamWriteValue64 / cuStreamWaitValue64 path.
• HIP backend: signal memory; native stream memops where the runtime records them under capture (ROCm 10), explicit batch-memop graph nodes where it does not (ROCm 7.14), decided by measuring once at startup. The probe does a real capture, instantiate and replay.
• When the handshake is unavailable and host-function nodes are not replay-safe (ROCm), graph capture raises and points to --cuda-graph-max-bs 0 instead of producing wrong outputs on replay. Refs [ROCm] CpuMoeExecutor decode is silently wrong under CUDA-graph replay (rel err 1.36, no exception) #350.

The HIP handshake and its probe are adapted from #378 and #667 by @zihaomu.

Kernel IO

• KernelIo::publish writes the step's x / ids / w into the pinned IO block through its mapped address, then writes done = 0 and the doorbell word. When none of the step's routes goes to the CPU it writes done = 1 and rings nothing, so the coordinator is not woken and the wait passes at once.
• KernelIo::combine follows the wait. It reads y from the mapped block and adds it into the GPU partial in place (hybrid) or writes it out (cpu backend).
• Together they replace the three D2H copies, the memop doorbell, the H2D copy and the torch add. The front-end memop wait stays. gpu_sync.h states the skip, and GpuSync gives KernelIo the device addresses of its two words.
• gpu_sync_io.cu is the only device source in _cpu_moe, which stays a CppExtension. On CUDA it is built as SASS for sm_75, sm_80, sm_90, sm_100 and sm_120 (each also runs on the newer minors of its major) plus compute_75 PTX that the driver builds for any other GPU. Naming the archs also keeps torch from picking them, which fails on a build machine without a GPU. On ROCm torch adds its offload architectures.
• At startup both kernels run once on scratch pinned memory (kernel_io_probe). If that fails, the executor logs a warning and keeps the memop doorbell with copies. ROCm builds the kernels but keeps the memop path until they are checked on AMD GPUs. The executor-ready log line names the handshake in use (kernel, memop or host-func).
• Kernel IO takes bf16 hidden states; a step with another dtype uses the copies.
• Coordinator: a task whose routes name no expert gets y zeroed and completes without waking the worker pool. The kernel path never rings for such a step, so this covers the copy path.

Behavior changes

• apply_router_weight_on_input with mxfp4_triton or ds_fp4 experts raises; both apply the router weight on the output. No current model sets it.
• top_k > 64 raises.
• A decode batch larger than max_tokens raises. The engine sizes max_tokens to cover explicit cuda_graph_bs as well.
• Extension API: the executor takes weight_format, banks and options; memop_submit / memop_sync, set_act_block / set_input_prequant, fp8_block_isa_name and max_weight_format_id are removed; enable_kernel_io, flag_publish, flag_combine and kernel_io_probe are added.

Validation

Refactor commit (51afda3)

Intel Xeon Platinum 8559C, RTX PRO 6000 Blackwell Server Edition, torch 2.14.1.

• Bit-exact against main (95ac86f): 384 executor cases covering 6 formats, 4 ISA tiers, the activations, apply-on-input, batch sizes 1/3/8, routing ids with duplicates and -1, and 2 layers. 0 mismatches in output and ISA name.
• tests/moe tests/engine tests/kernels/test_swiglu_clamp.py tests/models/deepseek_v41/test_engram.py: 276 passed, 6 skipped. The two fewer than main are the fp8_block avx512bf16 cases, which ran the same kernel as avx512. Graph-replay tests also pass with FREETOKEN_CPU_MOE_FLAG_SYNC=0.
• Decode throughput, bench_decode_moe.py --greedy --decode 512 --cache-rate 0.25, ABBA against main, tok/s:

Model	Strategy	main	this PR
Qwen3.6-35B-A3B FP8 (fp8_block)	cpu	81.07	80.76
	hybrid	125.38	125.19
Qwen3.6-35B-A3B NVFP4	cpu	103.44	105.46
	hybrid	143.66	144.79

Output tokens match main in every run.

• HIP: compiled and linked against the ROCm 7.14 SDK headers. Not run on AMD hardware. tests/moe/test_cpu_moe.py::test_cpu_moe_decode_cuda_graph_replay on ROCm 10 and on ROCm 7.14 would cover both handshake paths.

Kernel IO and coordinator skip (the two perf commits)

Intel Xeon Platinum 8480+, H100 80GB HBM3, torch 2.14.1+cu130.

• tests/moe tests/engine/test_moe_cpu_layers.py tests/kernels/test_swiglu_clamp.py: 157 passed, 6 skipped, the same as 51afda3 on this machine.
• Equivalence (a script, not added to tests/), one executor switched between kernel IO and the copies:
  • eager, batch 1 / 3 / 8, half or all routes on the CPU, with and without a GPU partial to add into: outputs bit-identical;
  • a step without CPU routes: kernel IO rings nothing and returns zeros or the unchanged GPU partial; the copy path rings once and the coordinator returns zeros;
  • a CUDA graph with 4 layers per replay that mixes ringing and skipped layers, replayed 12 times with fresh inputs: outputs equal to the eager copy path, doorbells served equal to the number of ringing layers, nothing pending afterwards.
• Loading: with CUDA_CACHE_DISABLE=1 CUDA_DISABLE_PTX_JIT=1 the executor decodes with handshake=kernel from the SASS; with CUDA_CACHE_DISABLE=1 CUDA_FORCE_PTX_JIT=1 it decodes with handshake=kernel from the PTX.
• Fallback, on a build with only the compute_75 PTX: with CUDA_CACHE_DISABLE=1 CUDA_DISABLE_PTX_JIT=1 the probe fails ("publish launch: PTX JIT compilation was disabled"), the executor logs the warning and decodes with handshake=memop.
• ROCm build: in the rocm/pytorch:rocm10.1.0_ubuntu24.04_py3.12_pytorch_release_2.14.0 image, setup.py build_ext with PYTORCH_ROCM_ARCH="gfx942;gfx1100;gfx1201" builds all three extensions, hipcc builds gpu_sync_io.cu for the three targets, and _cpu_moe imports. Not run on AMD hardware.
• Decode throughput, bench_decode_moe.py --greedy --decode 512 --cache-rate 0.25, ABBA against 51afda3, tok/s, measured on the build with only the compute_75 PTX (same sources; the SASS list was added after):

Model	Strategy	51afda3	this PR
Qwen3.6-35B-A3B-FP8 (fp8_block)	cpu	70.56 / 71.53	74.27 / 74.37
	hybrid	107.04 / 106.81	113.78 / 113.80
nvidia/Qwen3.6-35B-A3B-NVFP4	cpu	79.20 / 80.13	83.32 / 83.99
	hybrid	126.70 / 127.00	139.17 / 139.51

Output text sha1 matches 51afda3 in every run.

• Not covered: AMD hardware, Windows builds.

[jason-fxz]

Pushed two perf commits on top of the refactor:

• d975292 perf(cpu-moe): skip the worker pool for tasks without routed experts: the coordinator zeroes y and completes a task whose routes name no expert, without waking the worker pool.
• 886f752 perf(cpu-moe): publish and combine the cpu handshake with kernels: on CUDA, KernelIo (gpu_sync_io.cu) writes the step's inputs into the mapped pinned IO block and rings the doorbell from a kernel, or skips the doorbell when no route goes to the CPU; a second kernel reads the CPU rows back and adds them into the GPU partial. This replaces the three D2H copies, the memop doorbell, the H2D copy and the torch add. ROCm builds the same source but keeps the memop path for now.

Decode throughput on H100 80GB HBM3 and Xeon Platinum 8480+, bench_decode_moe.py --greedy --decode 512 --cache-rate 0.25, ABBA against the refactor commit 51afda3, tok/s:

Model	Strategy	51afda3	with the two commits
Qwen3.6-35B-A3B-FP8	cpu	70.56 / 71.53	74.27 / 74.37	+4.6%
	hybrid	107.04 / 106.81	113.78 / 113.80	+6.4%
nvidia/Qwen3.6-35B-A3B-NVFP4	cpu	79.20 / 80.13	83.32 / 83.99	+5.0%
	hybrid	126.70 / 127.00	139.17 / 139.51	+9.8%

Output text sha1 matches 51afda3 in every run. The numbers come from a build with only the compute_75 PTX, before the SASS list was added to setup.py; the sources are the same. The correctness checks, the fallback and the ROCm build are in the description.

[zihaomu]

I found a HIP probe failure-recovery issue while validating this PR on AMD Radeon AI PRO R9700 (gfx1201). I prepared a focused fix on top of the current PR head 886f752:

• Commit: d334eee — fix(cpu-moe): clean up failed hip capability probes
• Branch: fix/pr670-hip-probe-recovery-20261014

The branch adds exactly one commit after 886f752, limited to gpu_sync_hip.cpp and its fault-injection regression tests. The separate GPU alignment fix (7daf872) and DS-FP4 changes are excluded.

Failure and runtime scope

If hipGraphAddBatchMemOpNode succeeds but hipStreamUpdateCaptureDependencies fails, the newly added node remains outside the stream's dependency chain. I reproduced hipErrorStreamCaptureUnjoined (904) from EndCapture, with a residual HIP last-error. The probe returns unavailable and selects eager/host-function fallback, but the next ordinary PyTorch GPU operation can fail on the leftover error.

This is a fault-injection recovery issue, not evidence that normal ROCm 10 capture is broken. Normal ROCm 10 native WRITE/WAIT capture passed. ROCm 7.14 naturally selects the explicit-node compatibility path; reproducing the dependency-update failure on ROCm 10 requires forcing that compatibility path. I also checked the actual loaded HIP libraries and found no runtime-version mix-up.

Fix logic

1. Destroy the just-added graph node if updating capture dependencies fails, and return the original failure code. This allows capture cleanup without leaving an unjoined node.
2. Add an error scope to both native-memop detection and the handshake probe. On exit, after their existing resource cleanup, consume the HIP last-error generated during probing so eager fallback remains usable.
3. Check hipPeekAtLastError() before probing. If the caller already has an error, raise it without clearing it, rather than hiding it as an unsupported capability.

Validation

Single visible R9700; host CPU: dual AMD EPYC 9334. Separate rebuilt native extensions were used for both environments:

Runtime	PyTorch	HIP
ROCm 10.0.0	2.11.0+rocm10.0.0	7.15.26333
ROCm 7.14.0	2.11.0+rocm7.14.0	7.14.60850

The three new tests cover dependency-update failure, an EndCapture failure that sets a real HIP last-error, and a pre-existing caller error. They verify eager numerical results on changed inputs, balanced capture cleanup, node rollback, rejection of unsafe graph capture, and preservation of the caller's earlier error.

Targeted tests	ROCm 10	ROCm 7.14
Original 51afda3 production code/extension, with the new test harness	3 failed	3 failed
51afda3 + this HIP patch	3 passed	3 passed

The original combined-runtime fault/disable matrix improved from 19/22 to 22/22 passed. On the HIP-only repair of 51afda3, each runtime also passed 1,800 six-format changed-input graph replays, 1,000 replays across 25 live graphs, and 100 four-layer chain replays, all bit-exact against eager; 3,024 executor outputs per runtime matched the PR parent bit-exactly.

After updating to 886f752, I rebuilt the native extensions and repeated the three probe tests: 3 passed on each runtime. That latest-head run was on a broader local integration branch containing independent AMD fixes; I am not presenting it as a fresh full-model/performance validation of the linked single-commit branch, or as complete AMD sign-off.

To apply just this fix:

git fetch https://github.com/zihaomu/FreeToken.git fix/pr670-hip-probe-recovery-20261014
git cherry-pick FETCH_HEAD

The targeted regression entry point, after rebuilding the HIP extension in the configured environment, is:

python setup.py build_ext --inplace
python -m pytest -q tests/moe/test_cpu_moe.py -k test_cpu_moe_hip_probe_error_recovery

[jason-fxz]

Thanks @zihaomu, picked up d334eee as is: the branch now fast-forwards onto it.