fix(rocm): make CPU/Hybrid MoE graph replay safe - #378
Draft
zihaomu wants to merge 10 commits into
Draft
Conversation
- 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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes #350, where CPU/Hybrid MoE produces silently incorrect output during ROCm graph replay while eager execution remains correct.
hipMallocSignalMemoryand explicit graph batch-memory-op nodesCpuMoeExecutor, so their lifetime covers every graph replay without module-global growthmemop_submit/memop_syncpath to limit NVIDIA regression riskDependency 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:
cuStreamWriteValue64/cuStreamWaitValue64wrappersValidation
git diff --checkpassesDraft exit criteria