Skip to content

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

Draft
zihaomu wants to merge 10 commits into
FlashML-org:mainfrom
zihaomu:fix/issue-350-rocm-cpu-moe-graph
Draft

fix(rocm): make CPU/Hybrid MoE graph replay safe#378
zihaomu wants to merge 10 commits into
FlashML-org:mainfrom
zihaomu:fix/issue-350-rocm-cpu-moe-graph

Conversation

@zihaomu

@zihaomu zihaomu commented Sep 4, 2026

Copy link
Copy Markdown

Summary

Fixes #350, where CPU/Hybrid MoE produces silently incorrect output during ROCm graph replay while eager execution remains correct.

  • fail closed on ROCm: if the native stream-memory handshake does not pass a real capture/instantiate/replay probe, disable graph capture and continue on the correct eager path
  • add the ROCm 7.14 native handshake with hipMallocSignalMemory and explicit graph batch-memory-op nodes
  • keep HIP signal and graph-node parameter storage owned by CpuMoeExecutor, so their lifetime covers every graph replay without module-global growth
  • size per-layer flag capacity from the actual configured graph batch sizes, including configurations with more than 16 sizes
  • preserve the established CUDA module-level memop_submit/memop_sync path to limit NVIDIA regression risk
  • document the CUDA/ROCm synchronization difference and the ROCm fallback behavior

Dependency and scope

Depends on #132. This branch is based directly on the current #132 head (c0713e4). #132 itself is unchanged.

This is intentionally a Draft stacked PR while the native ROCm graph path receives RDNA3 and long-running model-serving stability validation. Until #132 merges, GitHub will also show the dependency commits in this PR; the diff will collapse to this single follow-up commit after the base lands.

CUDA and ROCm retain the same externally visible CPU-MoE ordering and results. Only their GPU/CPU synchronization implementations differ:

Backend Handshake implementation
CUDA Existing mapped-pinned per-slot arrays through module-level cuStreamWriteValue64/cuStreamWaitValue64 wrappers
ROCm 7.14 Executor-owned HIP signal memory plus explicit graph batch-memory-op nodes

Validation

  • PyTorch 2.11.0 + ROCm 7.14 on one Radeon AI PRO R9700: targeted safety and replay suite, 21 passed
  • capture/replay exercised across 19 graph batch sizes, covering the former fixed-16-slot boundary
  • eager plus repeated graph replay covered bf16, MXFP4, DS-FP4, and GGUF Q4_0 CPU-MoE formats
  • ROCm 7.2.4 / PyTorch 2.10 compatibility build passed and remains fail-closed where graph batch-memory-op APIs are unavailable
  • CUDA module API and dispatch routing are covered by unit tests; NVIDIA hardware validation remains required before this Draft is marked ready
  • git diff --check passes

Draft exit criteria

  • reproduce the replay suite on RDNA3
  • complete repeated graph rebuild/replay and long-running serving stability checks
  • obtain NVIDIA CI coverage for the preserved CUDA path

bouclem and others added 10 commits August 24, 2026 13:53
- Add hip_compat.h shim mapping CUDA runtime API to HIP equivalents
- Update pinned_tensor.cpp to compile under both nvcc and hipcc
- Add ROCm detection in arch.py (is_rocm, get_rocm_gfx_arch, is_gfx11xx_family)
- Guard NVIDIA arch checks to return None on ROCm
- Skip nvcc version check in _toolchain.py when on ROCm
- Add ROCm build path in setup.py (ROCM_HOME, amdhip64, --offload-arch)
- Add _hip_cflags() in kernel/utils.py for JIT compilation on ROCm
- Add is_rocm() and driver_hip_version() in backend.py
- Add rocm-smi fallback in __main__.py for clangd generation
- Add TODO(ROCm) for NCCL->RCCL, flashinfer/sgl_kernel ROCm builds,
  Triton autotune RDNA3 tuning, PDL equivalent, hiprtc JIT cache
- Add AMD ROCm classifier in pyproject.toml
Fail closed to eager execution when the HIP stream-memory handshake cannot survive capture and replay. Add a ROCm 7.14 graph batch-memop path with executor-owned signal and parameter storage, dynamically size graph flag slots, preserve the existing CUDA module API, and cover the safety and multi-format replay paths.
samuelishida pushed a commit to samuelishida/FreeToken that referenced this pull request Sep 5, 2026
What:
- Remove .agents/learnings and .plans/rocm-consolidation files from the branch.
- Remove internal increment and plan-path references from source comments and public installation docs.
- Keep implementation comments that explain correctness, ownership, profiler intent, source attribution, or ROCm safety behavior.
- Clarify public ROCm documentation: gfx1100 has recorded serving smoke on ROCm 7.2.1; the ROCm 7.14.x container is a reference environment, and other target cells remain compile-only until physical serving evidence exists.

Why:
- Keep merge surface focused on code, tests, reproducibility tooling, and user-facing documentation.
- Prevent private planning history, review workflow language, stale plan paths, and local process notes from entering the upstream repository.
- Avoid presenting compile success or a reference container as cross-target serving or performance proof.

Related upstream work informing this branch:
- PR FlashML-org#132: portable ROCm/HIP foundation.
- PR FlashML-org#133: TVM-FFI index/store portability.
- PR FlashML-org#135: RCCL tensor-parallel communication.
- PR FlashML-org#136: native GGUF ROCm build and Q4_0 kernels.
- PR FlashML-org#137: earlier AMD serving bring-up.
- PR FlashML-org#217: source-fork ROCm, Qwen3.5 GGUF, and performance experiments.
- PR FlashML-org#241: gfx1150 build, JIT, Triton, and attention hardening.
- PR FlashML-org#260: gfx1151 validation and fallback/build evidence.
- PR FlashML-org#316: HIP graph-capture-safe expert copies.
- PR FlashML-org#378: CPU/Hybrid MoE graph replay safety.
- Local branch milestones: 436263f, 926c1e8, e1d1856, 8a70c7e, and e5fd30f.

Evidence:
- 170 focused tests passed after cleanup.
- gfx1100 is the only target with end-to-end Qwen3.5 GGUF serving smoke recorded here.
- Remaining matrix targets are compile-only; no new throughput claim is published without a matching A/B manifest.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[ROCm] CpuMoeExecutor decode is silently wrong under CUDA-graph replay (rel err 1.36, no exception)

3 participants