Skip to content

fix(moe): prevent stale cpu expert inputs in rocm graphs - #764

Open
zihaomu wants to merge 1 commit into
FlashML-org:mainfrom
zihaomu:fix/rocm-cpu-moe-graph-io
Open

zihaomu wants to merge 1 commit into
FlashML-org:mainfrom
zihaomu:fix/rocm-cpu-moe-graph-io

Conversation

@zihaomu

@zihaomu zihaomu commented Oct 20, 2026 •

Copy link
Copy Markdown
Contributor

Fix stale CPU expert inputs during ROCm graph replay. On R9700, capture can succeed while replay returns incorrect layer outputs; eager decode is correct.

The fix reuses the existing mapped-memory publish/combine kernels on ROCm and checks multi-task graph outputs against eager before model capture. Decode graphs are disabled when the checks reject the IO path or dtype, with the decision shared across TP ranks. ROCm graph capture requires BF16 inputs.

Validation on one R9700 (gfx1201, ROCm 7.14.1, Torch 2.14.1+rocm7.14):

  • 72 tests passed, covering changing inputs/routes, skipped CPU routes, output accumulation and fallback. The stale-input regression fails on main and passes with this fix.
  • Qwen3-30B-A3B-FP8 completed 512/32K-input, 256-output requests using hybrid Graph. Median decode speeds across three fresh services were 92.78 / 61.66 token/s, respectively. These are validation results, not a same-mode Graph speedup comparison.

Hardware validation covers one R9700; TP agreement is tested with a mocked collective. Capture/runtime exceptions still surface as startup errors.

Related: #350, #670.

@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.

This branch has not been deployed

No deployments
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.

1 participant