fix(rocm): make CPU/Hybrid MoE graph replay safe

[zihaomu]

Summary

Fixes #350: CPU/Hybrid MoE can return numerically wrong, finite outputs during graph replay on ROCm even though eager execution is correct. This PR makes the CPU/GPU handshake replay-safe when supported and disables graph capture when it is not. This follow-up is now based directly on main 95ac86f21567bc244c7a9d91972838fbf049b4ed; #132 is merged, and #491's hybrid orchestration refactor is not included.

• Fail closed: probe the HIP handshake with real capture/instantiate/replay. If unavailable, disable Graph for CPU/hybrid and offload with CPU layers, retain correct eager execution, and reject unsafe direct executor capture.
• Use HIP signal memory and explicit graph batch-memory-op nodes. Keep signals and immutable per-slot parameter storage executor-owned, including across repeated graph rebuilds.
• Derive capacity from the configured graph sizes, including more than 16 sizes and explicit sizes above the normal decode limit.
• Agree on safety across the TP CPU process group: one failed probe disables capture on every rank.
• Clean up partially failed probe capture, and keep the timeout release alive until a delayed probe launch drains.
• Preserve CUDA's established module-level memop entry points and current-main quantization/CPU-threading behavior. Native ROCm graph support is enabled by the probe, never the version number alone.

Validation

Author-run on one AMD Radeon AI PRO R9700 (gfx1201, 32 GiB), dual EPYC 9334 host, with test processes pinned to CPUs 32-39. Only one GPU was exposed and torch.cuda.device_count() == 1 was checked. Reference environment: PyTorch 2.11.0+rocm7.14.0, HIP 7.14.60850, Triton 3.7.1, Python 3.12. Synthetic banks only; no model checkpoint, quality, serving-throughput, or full-suite claim.

python setup.py build_ext --inplace
PYTHONPATH=python taskset -c 32-39 python -m pytest \
  tests/engine/test_cpu_moe_graph_safety.py \
  tests/engine/test_moe_cpu_layers.py \
  tests/moe/test_cpu_moe.py \
  tests/moe/test_cpu_moe_q4_0.py \
  tests/moe/test_cpu_moe_fp8_block.py \
  tests/moe/test_dsfp4_prequant.py \
  tests/moe/test_hybrid_fetch.py \
  tests/moe/test_offload.py -q --tb=short -m 'not slow' \
  -k 'not test_adjust_config_converts_moe_cache_rate_to_cache_size'

113 passed, 2 deselected. The exclusions are a slow GGML comparison and an existing test hard-coding CUDA FlashInfer; the latter was independently checked to fail identically on clean main and this branch on ROCm.

• Covers BF16, MXFP4, DS-FP4, Q4_0 and block-FP8 replay; real hybrid CPU/GPU computation with overlap on/off; numerical comparisons across repeated rebuilds and 19 simultaneously live graphs.
• Core safety/CPU-MoE files with FREETOKEN_CPU_MOE_FLAG_SYNC=0: 63 passed, 13 skipped, 1 slow deselected. Unsupported native Graph cases skip; direct-capture safety and eager correctness remain tested.
• CPU-only safety, CPU-layer resolution and block-FP8 files: 41 passed, 10 skipped.
• Supplemental local stress: eight independent processes, 56 numerical/rebuild cases, 696 graph replays. This is not a model-serving soak.
• Local fault injection: unsupported graph-node API and launch delayed beyond the probe deadline both return unsafe and preserve correct eager results. The unsupported-API cleanup regression failed before and passed after the cleanup fix.
• Capacity regression tests fail on clean main and pass here. Startup safety consensus is tested with two real CPU/Gloo ranks, not two GPUs.
• All native extensions also compile under the local ROCm 10 image, and _cpu_moe imports successfully without a GPU. No ROCm 10 GPU replay claim.
• git diff --check passes.

Remaining validation

The original RDNA3 and NVIDIA physical-hardware checks and long-running CPU/Hybrid model-serving criteria are not yet satisfied by this author's evidence. Full TP>1 GPU replay remains untested. The AMD reference image uses Torch 2.11; the new main packaging defaults to Torch 2.14 and that default dependency environment was not validated here. These remain open validation items while this PR is ready for review.

Concrete ROCm counterexample (rerun 2026-10-13)

The required behavior is simple: capture a CPU MoE decode with input A, then replace the hidden states, expert IDs, and router weights in the same GPU tensors with input B and replay. The replay must compute B's result, matching eager execution of B.

CPU MoE needs an ordered sequence on every replay:

copy current input/routing GPU -> CPU
  -> notify the CPU expert workers
  -> wait for the CPU result
  -> copy the result CPU -> GPU

In the tested ROCm environment, main's native flag probe returns false and selects the host-function fallback. That fallback does not preserve this handshake correctly during graph replay. Replay and GPU synchronization complete without a runtime error, and the output is finite, but numerical comparison with the current-input reference fails. Graph API success alone therefore cannot establish correctness.

I reran main 95ac86f21567bc244c7a9d91972838fbf049b4ed and this PR de45a143bbf11ec4728ba0235e6bc3ad55f76bf1 on the same single R9700 and PyTorch/HIP environment listed above. A supplemental local BF16 script used identical banks and inputs for both checkouts, with an independent Torch reference for expert computation. It uses one layer, 8 experts, H=512, I=256, batch=4, and top_k=2.

New-input replay	main eager error	main Graph error	PR eager error	PR Graph error
0	0.002539	0.574833	0.002539	0.002539
1	0.002068	0.282794	0.002068	0.002068
2	0.002419	0.484713	0.002419	0.002419

Error is max(abs(actual - reference)) / max(abs(reference)); the BF16 acceptance threshold is 0.02. Main fails 3/3 replays, while the PR passes 3/3 and its Graph errors equal its eager errors. These results establish incorrect replay; they do not imply that every incorrect output equals a particular capture-time or stale buffer.

The supplemental script is local, not a committed test file. Reviewers can reproduce the same failure class with the repository's existing tests, without that script. In a separate fresh run:

Existing graph test	main replay-0 error	Threshold	main	PR
BF16	0.618253	0.02	FAIL	PASS
MXFP4	0.346774	0.04	FAIL	PASS
DS-FP4	0.811441	0.04	FAIL	PASS
Q4_0	0.783228	0.02	FAIL	PASS

Main: 4 failed. PR: 4 passed, with no skips. Exact failing errors may vary with execution. The smallest reproducer, run separately in each checkout after rebuilding its extensions, is:

python setup.py build_ext --inplace
PYTHONPATH=python taskset -c 32-39 python -m pytest \
  tests/moe/test_cpu_moe.py::test_cpu_moe_decode_cuda_graph_replay \
  -q --tb=short

CPU affinity 32-39 is specific to this host; select suitable cores elsewhere.

The fix makes the CPU/GPU notification and wait explicit HIP Graph nodes backed by signal memory, and enables them only after a real capture/instantiate/replay probe succeeds. In the successful probe path, the above numerical replay cases pass. If the handshake is unavailable, Engine disables Graph and executes eagerly; direct executor capture raises an explicit error. A separate run with FREETOKEN_CPU_MOE_FLAG_SYNC=0 confirmed correct eager output and the intended capture rejection.

This counterexample concerns CPU/Hybrid MoE synchronization, including offload configurations with CPU layers. ROCm's general Graph support and pure-GPU computation are outside this failure. The measured benefit here is correctness; no end-to-end throughput improvement is claimed.

[lukascechovic]

Thanks for building this — we filed #350 and did not expect a fix, let alone one validated on the
same card we found it on. Your framing of the fix as fail-closed — disable capture and
continue on the correct eager path when the native handshake does not pass a real
capture/instantiate/replay probe — is the property we would have asked for and did not know how to
ask for. A silent wrong answer converted into a slower right one is exactly the right trade for
this class of bug.

Two things below: a ROCm datapoint your HIP memop work bears on directly, and an offer.

1. Your HIP memop port has a second consumer, and it is the half that is still private

Reading the diff, the thing that struck us is that the HIP hipStreamWriteValue64 /
hipStreamWaitValue64 / hipStreamBatchMemOp port is already here — we had it filed on our
own side as a piece of work someone would eventually have to do. It matters to us because the CPU
MoE executor is not the only thing in this engine that wants that handshake.

#311's disk-streamed PLE table resolves its stream memops by dlopening libcuda.so.1. On
ROCm that dlopen fails, the probe returns false, and the table falls back to launch-gating: a
blocking D2H of the sampled token that drains the stream, then hash, io_uring read, H2D and launch
on the host, with the GPU idle across the whole stretch. Measured on our box — half-box (one
R9700, TP=1), on tree c4b96fc1 (upstream 4b94bdc3 plus our patches 0001..0009), with
--ple-backend the only flag changed between the arms:

arm	decode at 10k	decode at 100k
--ple-backend pinned	26.100 ± 0.351 tok/s (n=8/10)	26.345 ± 0.110 (n=2/3)
--ple-backend disk	18.984 ± 0.272 tok/s (n=9/10)	19.432 ± 0.182 (n=2/3)

n is post-filter: replies are dropped if their zlib compression ratio falls below 0.40, which is
how we catch a degenerate repeated-token generation before it flatters a decode number. Note the
100k pair rests on two samples per arm — we quote it because it agrees with the 10k pair, not
as an independent result.

−27.3 % at 10k, −26.2 % at 100k, roughly 14 ms per token, against the −2.6 % … +0.2 %
band #311 reports on H100 and RTX PRO 6000 Blackwell — which was measured on the wait-sync
path. Prefill is unaffected: TTFT +0.20 s at 10k and −0.22 s at 100k, same half-box arms —
though every prefill number from this box carries "gfx1201, BLOCK_N clamped to 32" (see
Platform), and block_n = 64 does not run here at all — it raises OutOfResources before
executing an instruction — so the clamp's cost is unmeasurable rather than zero. The memory win is
real and large — on the same arms the PLE table stops pinning 47.7 GiB of host RAM — which is
why we run it anyway.

We can only bound the split, not measure it. Sampled from /proc/diskstats on the same
half-box disk arm, during decode the NVMe served 300 reads/s at 0.648 ms average = 16 reads per
token, one 4K O_DIRECT read per row. Fully serialised that is 10.4 ms; if the store's io_uring
batch puts all 16 in flight together it is ~0.65 ms. /proc/diskstats cannot tell those apart, so
the disk accounts for 0.65–10.4 ms and the launch-gating stall for the remaining 4–13.7
ms. What the bound does establish is that the disk alone cannot reach 14 ms even in its worst
case. Splitting it properly needs a wall-clock inside the fill, which we have not written — so
treat the split as bounded, not measured.

Where your PR lands on it: the eager half is now reachable — on ROCm your module-level
memops_probe / memop_submit / memop_sync resolve to the HIP implementations, so a second
consumer could stop dlopening libcuda and just call them. The half that is still
CpuMoeExecutor-owned is the capture half — and that is the half #311 needs, since by its own
description the decode graph launches first and waits on the flag right before consuming the rows.
If the ROCm graph-node splicing ever becomes shareable, there is a caller waiting for it.

And the ROCm-7.14 fact you documented in that file — that ordinary hipStreamWrite/WaitValue
and hipStreamBatchMemOp calls execute during capture but are not recorded in the graph, so
you add explicit batch-memory-op nodes and splice them into the dependency set — is worth flagging
loudly to anyone attempting a second port. Someone who ports the eager calls, sees them work, and
captures a graph gets silently stale values on replay: the same failure shape as #350 itself, in a
different file.

None of this is a request. It is a note that the two land in the same place.

2. An offer against your draft-exit criteria

You list "reproduce the replay suite on RDNA3", "complete repeated graph rebuild/replay and
long-running serving stability checks", and NVIDIA CI for the preserved CUDA path.

We can help with parts of the middle one, on RDNA4 rather than RDNA3. We run two Radeon AI
PRO R9700 (gfx1201) on Linux under ROCm 7.14 / PyTorch 2.11, serving a 100 GiB-class MoE
continuously, and we can run your targeted safety and replay suite
(tests/engine/test_cpu_moe_graph_safety.py and the tests/moe/ graph-replay cases) on both
cards and report the numbers the way we did on #132.

We can also exercise it at TP=2. To be exact about what is new there: we have reported TP=2
serving before, on #135, so the gap is not TP=2 as such — it is that nobody has run the
CPU/Hybrid graph-replay path at TP>1, on any card. That is worth closing, because from our own
experience a caught per-rank exception behaves very differently at TP=2 than at TP=1. That
difference is a scheduler problem rather than yours, and we have written it up separately.

Three honest caveats before you count on any of it.

We cannot give you the "long-running serving stability" half. We serve --moe-backend offload, not cpu or hybrid, so a serving soak on this code path is not something our
production traffic would ever produce — it would be a rig we stand up specially, and it would not
be representative of a real workload the way our offload row is. Repeated rebuild/replay and the
suite are the parts we can do well.

We cannot promise a date. The box is a shared machine with a live soak on it. Say what you
want run and we will fit it in.

It is a second gfx1201 datapoint, not a clean-main CI result. RDNA4, on a tree that is
upstream 4b94bdc3 plus a local patch series, with the BLOCK_N clamp in the Platform note
below.

No quality claim from us in any case — we have no fidelity instrument on this box, which is why
#350 was reported as a numerical-deviation finding with a relative error rather than as a quality
regression.

---

Platform

2 × AMD Radeon AI PRO R9700 (gfx1201, 32 GiB each), Linux, ROCm 7.14, PyTorch 2.11, 6 physical
cores / 122 GiB host RAM. Serving qwen4_exp / Qwen3.8-Flash-Next NVFP4 with --moe-backend offload; the CPU/Hybrid finding in #350 came from running this repo's kernel suite on gfx1201,
not from this serving configuration.

Our tree is upstream 4b94bdc3 plus a local patch series, whose build asserts tree c4b96fc1
for the 0001..0009 stack every §1 number was measured on. The QSA prefill tile is halved until 2 · head_dim · block_n · itemsize fits a 60 KiB LDS budget (#349) — at this checkpoint's head_dim = 256 in bf16 that lands on BLOCK_N = 32. So every prefill number here is on the clamped tile
and is not comparable to a CUDA run of the same engine.

Box labelling, since this box has two cards: every number in §1 — decode, TTFT, the 47.7 GiB
and the /proc/diskstats figures alike — is half-box, one R9700 at TP=1, on tree c4b96fc1
(upstream 4b94bdc3 plus our patches 0001..0009), with --ple-backend the only flag changed
between the arms. Anything we quote at --tp-size 2 is full-box. Every number is named
against the tree it was measured on, and none of them is measured on upstream main.

[zihaomu]

Rebased this PR onto the refreshed AMD foundation (#132, head c2af4c6) and updated it for the current moe_strategy / quantization interfaces.

Conflict resolution kept the current upstream Engine/CpuMoeExecutor contracts (sample.alpha, sample.limit, quant_method.cpu_format) while layering the ROCm graph-safe HIP signal/memop path on top.

Validation:

• tests/engine/test_cpu_moe_graph_safety.py: 14 passed
• CPU MoE suites: 28 GPU-dependent cases collected/skipped in the no-GPU test container
• all native extensions, including _cpu_moe, compile with ROCm 7.14 modular SDK
• git diff --check passes

New head: 94f77248886029c8c0072960c3d00a901b60deaa.