Skip to content

fix(moe): publish rocm cpu expert inputs before notification - #764

Merged
jason-fxz merged 4 commits into
FlashML-org:mainfrom
zihaomu:fix/rocm-cpu-moe-graph-io
Nov 3, 2026
Merged

jason-fxz merged 4 commits into
FlashML-org:mainfrom
zihaomu:fix/rocm-cpu-moe-graph-io

Conversation

@zihaomu

@zihaomu zihaomu commented Oct 20, 2026 •

Copy link
Copy Markdown
Contributor

CPU expert graph replay can observe a ready notification before its D2H inputs are visible on affected ROCm stacks, silently using stale activations or routes. Qualify the IO path before model capture and establish input publication before notifying the CPU.

  • Prefer KernelIO on ROCm under the existing CUDA routing conditions. When kernel_io_probe passes, BF16 inputs use the existing mapped-memory publish/combine kernels. The publish kernel establishes system ordering and skips CPU work for empty routes; this path needs no publication event. Shared arithmetic kernels are unchanged.
  • Keep an ordered copy fallback. When KernelIO is unavailable or inputs are FP16/FP32, use copies. With the memop handshake, record hipEventReleaseToSystem with hipEventRecordExternal between D2H copies and the ready notification, so every Graph replay retains D2H → event → ready. Events belong to each (layer, batch) task and remain alive through replay/recapture. Detect native memop capture at runtime and add explicit graph nodes when needed. Eager execution uses host callbacks if the memop handshake is unavailable. No custom HIP library is required.
  • Check the model's activation dtype against synchronous CPU execution. The startup numerical check covers eager and Graph outputs, changing inputs/routes, skipped CPU tasks and Hybrid output accumulation. A BF16 KernelIO numerical failure retries event + memcpy and checks it again. If eager is correct but Graph is unsafe, warn and disable decode graphs on every TP rank, clearing both explicit batch sizes and the maximum batch size. A failed eager check or a capture/runtime API error stops startup. When graphs are configured off, run only the eager numerical check. BF16 can use KernelIO; FP16/FP32 use copies and are qualified for their own dtype.
  • Enable ROCm locked hybrid staging. OS-locked host expert banks can use pinned staging with the same --moe-locked-fetch / --moe-locked-max-fetch eligibility as CUDA. The gather waits through stream memops when available or a GPU flag wait otherwise; the latter can construct GpuSync without requiring memops. Update AMD installation guidance for the IO selection, startup checks, fallbacks and native rebuild.

Validation reported by @LZ-QWQ in the latest test/model results:

  • R9700 and W7900D each pass 146 + 4 selected repository tests, TP=2 checks and exact event/KernelIO model logits/token comparisons. Without the event, the diagnostic copy path fails 10/10 numerical runs on R9700; event and KernelIO each pass 10/10 on both GPUs. W7900D did not reproduce the no-event failure.
  • RTX4090 and RTX5090 each pass 126 + 4 repository tests, with 20 HIP-only skips, plus Engine configuration, KernelIO/copy/host-function checks and TP=2 fallback checks.
  • Coverage includes the synchronous-CPU oracle, eager/Graph numerical rejection, activation dtypes, Hybrid accumulation, empty routes, replay/recapture, event API errors, and locked-fetch staging with both wait modes. Restoring the HIP exclusion makes the locked-fetch staging tests fail; the updated routing passes them.

Layer latency, FP8 weights/BF16 activations, H2048/I768/E8/K4, three layers, batch 1, four CPU workers, Hybrid accumulation and one CPU expert per token; medians across three processes:

GPU Event + copies + memops (µs) KernelIO (µs) Reduction vs event
R9700 103.08 91.92 10.83%
W7900D 112.56 94.74 15.83%

Locked-fetch decode, KernelIO/BF16, batch-1 Graph, 512 input / 128 greedy output tokens, Triton attention, six CPU workers, 8 GiB expert cache, moe_cpu_layers=2, hybrid_max_fetch=1; enabled settings are moe_locked_fetch=0.5, moe_locked_max_fetch=2:

GPU / FP8 model Locked off (token/s) Locked on (token/s) Change
R9700 / Qwen3.6-35B-A3B 31.634 35.473 +12.13%
W7900D / Qwen3.5-35B-A3B 41.579 43.177 +3.84%

Medians of five samples, except W7900D locked-on uses four after excluding one contended sample; W7900D off/on ranges are 41.557–41.646 / 42.779–43.394 token/s. These compare the candidate's IO paths/settings, not whole-main throughput. Full accuracy evaluation, long-duration stability and other ASICs remain outside the reported validation. The runtime follow-up supports retaining explicit publication on both tested AMD GPUs without claiming a general HIP notification guarantee or a deployed CLR fix.

Rebuild the native extension in the matching ROCm environment before serving:

python setup.py build_ext --inplace

Relevant repository checks:

PYTHONPATH=python python -m pytest \
  tests/moe/test_cpu_moe.py tests/moe/test_cpu_moe_fp8_block.py \
  tests/engine/test_moe_cpu_layers.py tests/moe/test_mixed_expert_formats.py \
  tests/moe/test_hybrid_fetch.py tests/moe/test_dsfp4_prequant.py \
  tests/moe/test_offload.py -q

Related: #350, #670; runtime investigation/fix: ROCm/rocm-systems#13877.

@github-actions github-actions Bot added the amd label Oct 20, 2026
KarrAcaRn pushed a commit to KarrAcaRn/FreeToken-ByAI that referenced this pull request Oct 20, 2026
FlashML-org#764 FlashML-org#763 FlashML-org#708 FlashML-org#687 deferred, FlashML-org#688 kept as a feature; FlashML-org#169 kept after its upstream close,
FlashML-org#494/FlashML-org#154/FlashML-org#315 replaced upstream by FlashML-org#690/FlashML-org#698; FlashML-org#278 dropped with upstream's auto->hybrid pick (FlashML-org#734).

Assisted-by: Claude Opus 5.5
@zihaomu

zihaomu commented Oct 21, 2026

Copy link
Copy Markdown
Contributor Author

Runtime-level fix: ROCm/rocm-systems#13877.

The independent HIP reproducer shows that the host can observe a stream WRITE notification before data copied earlier on the same stream becomes visible. The CLR fix adds a system-scope release before the notification and captures that barrier for every graph replay.

This FreeToken PR provides the application-level workaround and startup safety checks: use mapped-memory publish/combine kernels and reject unsafe graph paths. It remains useful on unpatched runtimes.

The ROCm PR is still under review and targets develop; release backports are separate. The model results above validate this FreeToken workaround on ROCm 7.14.1, not the patched CLR. Before retiring the workaround, we still need to rerun FreeToken's original memcpy + memop path with the patched runtime, including the multi-layer regression and Qwen3-30B-A3B-FP8 hybrid Graph validation.

@jason-fxz

Copy link
Copy Markdown
Collaborator

Thanks for tracking this down to the runtime. One thing before merge: docs/install_amd.md is out of date. Its "CPU and hybrid MoE" section still says that when the startup probe fails, CUDA graph capture stops with an error and to run with --cuda-graph-max-bs 0. With this PR the engine turns off decode graphs itself and logs a warning, and ROCm graphs also require the kernel IO path, BF16 inputs and a passing _probe_kernel_graph_io check. Could you update that section in this PR?

CUDA check on one H100 80GB (driver 580.95.05, torch 2.14.1+cu130):

PYTHONPATH=python pytest tests/moe/test_cpu_moe.py tests/moe/test_cpu_moe_fp8_block.py tests/engine/test_moe_cpu_layers.py
  • This PR: 57 passed, 6 skipped. The 3 extra skips are the new ROCm-only cases.
  • Base commit cc56b6b: 50 passed, 3 skipped.
  • test_cpu_fp8_block_graph_distinct_layer_inputs also passes against the base commit's code.

Not tested: ROCm.

@zihaomu zihaomu changed the title fix(moe): prevent stale cpu expert inputs in rocm graphs fix(moe): publish rocm cpu expert inputs before notification Oct 27, 2026
@LZ-QWQ

LZ-QWQ commented Oct 30, 2026

Copy link
Copy Markdown
Contributor

This update enables BF16 KernelIO and locked-fetch on ROCm under CUDA's routing conditions, without changing shared arithmetic kernels. Startup validates the model dtype against synchronous CPU execution, with ordered-copy fallback and TP-consistent Graph-to-eager fallback; eager/API errors remain fatal.

The failure is consistent with ready becoming visible before D2H inputs. CUDA documents a default pre-write barrier similar to __threadfence_system(); HIP's stream-write documentation does not state an equivalent guarantee. GpuSync::gpu_submit therefore records a system-release event (hipEventReleaseToSystem, hipEventRecordWithFlags(..., hipEventRecordExternal)) before hipStreamWriteValue64(..., hipStreamWriteValueDefault) or Graph memop publication: D2H → event → ready. KernelIO retains its existing system fence and avoids this event.

Validation:

  • Without the event, R9700 fails 10/10 numerical runs; event and KernelIO each pass 10/10 on both R9700 and W7900D. W7900D did not reproduce the no-event failure.
  • Each AMD GPU passes 146 + 4 selected repository tests, TP=2 checks, and exact event/KernelIO model logits/token comparisons. Locked-fetch staging tests fail with the HIP exclusion restored and pass with this change, including both wait modes and Graph replay/rebuild.
  • RTX4090 and RTX5090 each pass 126 + 4 repository tests (20 HIP-only skips), plus Engine configuration and KernelIO/copy/host-function checks, including TP=2 fallback.

Layer latency (µs): FP8/BF16, H2048/I768/E8/K4, three layers, batch 1, four CPU workers, Hybrid accumulation, one CPU expert/token; medians across three processes.

GPU Copies+memops, no event (unsafe) Event+copies+memops KernelIO Reduction vs event
R9700 95.70 103.08 91.92 10.83%
W7900D 112.05 112.56 94.74 15.83%

No-event timings are diagnostic. W7900D's small event/no-event difference is within process variation.

Locked-fetch decode (tokens/s): KernelIO, BF16, batch-1 Graph, 512 input/128 greedy output tokens, Triton, six CPU workers, 8 GiB expert cache, moe_cpu_layers=2, hybrid_max_fetch=1. Enabled: moe_locked_fetch=0.5, moe_locked_max_fetch=2.

GPU / FP8 model Locked off Locked on Change
R9700 / Qwen3.6-35B-A3B 31.634 35.473 +12.13%
W7900D / Qwen3.5-35B-A3B 41.579 43.177 +3.84%

Medians of five samples; W7900D locked-on uses four after excluding one contended sample. Its off/on ranges are 41.557–41.646 / 42.779–43.394 tokens/s. These compare candidate paths/settings, not whole-main TPS.

@LZ-QWQ
LZ-QWQ force-pushed the fix/rocm-cpu-moe-graph-io branch from ac7fb19 to 5e1f3b3 Compare October 30, 2026 10:46
@LZ-QWQ

LZ-QWQ commented Oct 30, 2026

Copy link
Copy Markdown
Contributor

With identical loaded HIP/ROCr/COMGR binaries and default allocations, our no-event tests observed ready before payload visibility on R9700; we did not reproduce this on W7900D. In the traced copy path, the payload copy ends with an AGENT-scope release, which alone does not guarantee CPU visibility. Inserting a system-release event between the copy and hipStreamWriteValue32/64 resolved the R9700 failures by making the payload visible to the CPU before signaling ready.

Our controlled tests and firmware analysis suggest two possible reasons W7900D did not reproduce the issue: its UC (uncached) payload mapping, and extra GPU L2 cache writeback (GL2_WB) in the firmware's SYSTEM-acquire handling before ready is written. The corresponding R9700 path appears to lack that extra writeback. This is consistent with W7900D passing after forcing the payload to NC (non-coherent) in eager mode, but failing in Graph with AGENT acquire. These findings concern the tested stack, not a general publication guarantee.

CUDA's default StreamWriteValue contract includes a pre-write fence similar to __threadfence_system(), scoped to the stream; HIP's documentation describes execution ordering without explicitly promising the same fence. For this HIP copy/memop path, retaining payload → system-release external event → ready on both GPUs is therefore appropriate: hipEventReleaseToSystem provides explicit publication, and hipEventRecordExternal preserves the event node during Graph capture. In the measured W7900D one-CPU-expert workload, median latency was 112.05 → 112.56 µs (+0.46%), within process variation; no clear regression was observed in this workload. Event overhead remains workload-dependent.

@jason-fxz

jason-fxz commented Nov 3, 2026 •

Copy link
Copy Markdown
Collaborator

Thanks for the update. Before merging, could you refresh the PR description? It still describes the 10-27 version, and it will become the squash commit message. Parts that no longer match the head (5e1f3b3):

  • It says ROCm needs no publish/combine kernel and that only CUDA keeps kernel IO. ROCm now uses kernel IO first when kernel_io_probe passes, and switches to event + memcpy only when the startup check fails with BF16 kernel IO.
  • The startup check now compares against synchronous CPU execution, including hybrid accumulation, instead of an eager reference.
  • "Numerical rejection disables decode graphs with a warning" now covers only a graph failure. A failed eager check is a startup error.
  • The graph gate is the model's activation dtype, not BF16 only.
  • The ROCm locked hybrid staging change in 5e1f3b3 is not mentioned.
  • The validation section and the overhead table treat the event path as the chosen one. The R9700/W7900D results and the kernel IO vs event layer latency from your later comments would fit better there.

@jason-fxz
jason-fxz merged commit 9cacd08 into FlashML-org:main Nov 3, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants