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

docs/install_amd.md

@@ -49,7 +49,36 @@ or NVFP4 Marlin/b12x (e.g. `--quant-backend moe.nvfp4=marlin` or
 ## CPU and hybrid MoE
 
 `--moe-strategy cpu`, `hybrid` and `--moe-cpu-layers` sync the GPU with the CPU
-expert pool through HIP stream memory operations. ROCm 10 records them into CUDA
-graphs itself; on ROCm 7.14, which does not, FreeToken adds the graph nodes
-explicitly. When the startup probe of that path fails, CUDA graph capture stops
-with an error; run those strategies with `--cuda-graph-max-bs 0`.
+expert pool through HIP stream memory operations. As on CUDA, BF16 inputs use
+kernel IO when its startup probe passes: a publish kernel copies the inputs to
+mapped host memory and notifies the CPU; a combine kernel reads the CPU output.
+The publish kernel establishes the memory ordering before notification and
+skips CPU work when no route needs it.
+
+When kernel IO is unavailable or the input is not BF16, FreeToken uses copies.
+With the memop handshake, this fallback records a system-release event after
+the input copies and before notifying the CPU. The event is recorded as an
+external graph node so that every replay retains the publication point. The
+kernel IO path does not record this event. If the memop handshake is unavailable,
+eager execution uses host callbacks instead. No custom HIP library is required.
+
+Native memop capture is detected at runtime. If it is unavailable, FreeToken
+adds explicit memop graph nodes. Before model capture, a numerical check
+compares eager and Graph outputs against synchronous CPU execution, including
+changing routes, skipped tasks and Hybrid output accumulation. If kernel IO
+fails this check, FreeToken switches to the event + memcpy path and checks it
+again. Startup checks only the model's activation dtype; BF16 can use kernel IO,
+while FP16 and FP32 use copies. When decode graphs are disabled, the numerical
+check runs eagerly without capturing a graph. If the Graph path is unsafe,
+startup logs a warning and disables decode graphs for the engine. TP ranks
+agree on this decision, so one unsafe rank disables graphs for all ranks.
+A failed eager check, or a capture/runtime API error, stops startup rather
+than allowing unverified output.
+
+Hybrid layers whose host banks are OS-locked can also fetch experts through
+pinned staging, using the same `--moe-locked-fetch` and
+`--moe-locked-max-fetch` settings as CUDA. The staging gather waits through
+stream memops when available, or a GPU flag wait otherwise.
+
+After updating the native CPU MoE sources, rebuild the extension with
+`python setup.py build_ext --inplace` in the ROCm environment before serving.

python/freetoken/engine/engine.py

@@ -43,6 +43,17 @@
 
 logger = init_logger(__name__)
 
+def _disable_unsafe_cpu_moe_graphs(config: EngineConfig, executor, dtype, tp_cpu_group) -> None:
+    safe = torch.tensor([executor.graph_capture_safe(dtype)], dtype=torch.int32)
+    if config.tp_info.size > 1:
+        torch.distributed.all_reduce(safe, op=torch.distributed.ReduceOp.MIN, group=tp_cpu_group)
+    if not safe.item():
+        logger.warning_rank0("CPU MoE IO/handshake is not graph-safe on every TP rank for this dtype; disabling decode graphs")
+        # Clear explicit sizes too; GraphRunner gives them priority over max_bs, including on rebuild.
+        object.__setattr__(config, "cuda_graph_bs", [])
+        object.__setattr__(config, "cuda_graph_max_bs", 0)
+
+
 def _clamp_decode_batch_to_pools(config: EngineConfig, pool_sizes: list[int]) -> None:
     """Clamp the decode batch so every GPU-decoded layer's routed experts fit in its pool."""
     num_experts = config.model_config.num_experts
@@ -1077,6 +1088,11 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None:
         # round a batch up to the largest captured size; cover both.
         max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, *(config.cuda_graph_bs or []), 1)
         max_tokens *= config.decode_rows_per_req
+        # Match GraphRunner's explicit-size precedence, including an empty list.
+        if config.cuda_graph_bs is not None:
+            graph_enabled = bool(config.cuda_graph_bs)
+        else:
+            graph_enabled = config.cuda_graph_max_bs is None or config.cuda_graph_max_bs > 0
         by_id = sorted(layers, key=lambda layer: layer.layer_id)
         layer_formats = [layer.quant_method.cpu_spec() for layer in by_id]
         executor = CpuMoeExecutor(
@@ -1087,12 +1103,16 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None:
             num_threads=config.moe_cpu_threads,
             max_tokens=max_tokens,
             device=self.device,
+            activation_dtype=self.dtype,
+            graph_enabled=graph_enabled,
             swiglu_alpha=float(sample.alpha),
             swiglu_limit=sample.limit,
             layer_formats=layer_formats,
         )
         cache.set_cpu_executor(executor)
         self.cpu_moe_executor = executor
+        if graph_enabled:
+            _disable_unsafe_cpu_moe_graphs(config, executor, self.dtype, self.tp_cpu_group)
 
     def _sync_get_memory(self) -> Tuple[int, int]:
         """Get the min and max free memory across TP ranks."""

python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp

@@ -157,13 +157,14 @@ hipError_t wait_done(hipStream_t stream, uint64_t* done, hipStreamBatchMemOpPara
 struct GpuSync::Impl {
   uint64_t* ready = nullptr;  // signal memory
   uint64_t* done = nullptr;
+  std::map<uint64_t, hipEvent_t> publication_events;
   // Explicit nodes keep these addresses, so one parameter pair per distinct task, never moved (std::map nodes are stable).
   std::map<uint64_t, std::array<hipStreamBatchMemOpParams, 2>> submit_ops;  // done = 0, ready = task
   hipStreamBatchMemOpParams wait_op{};
 };
 
-// The memop-free mode serves only the locked fetch, which ROCm keeps off: every GpuSync here needs memops.
-GpuSync::GpuSync(bool) : impl_(std::make_unique<Impl>()) {
+// Locked fetch can use the signal words with a kernel wait when memops are unavailable.
+GpuSync::GpuSync(bool require_memops) : impl_(std::make_unique<Impl>()) {
   hipError_t rc = alloc_signal(&impl_->ready);
   if (rc == hipSuccess) rc = alloc_signal(&impl_->done);
   if (rc != hipSuccess) {
@@ -172,10 +173,12 @@ GpuSync::GpuSync(bool) : impl_(std::make_unique<Impl>()) {
     hip_check(rc, "signal allocation");
   }
   // Measure now: the measurement captures and allocates, which a graph capture in progress would reject.
-  (void)records_native_memops();
+  if (require_memops) (void)records_native_memops();
 }
 
 GpuSync::~GpuSync() {
+  for (const auto& [task, event] : impl_->publication_events)
+    if (event != nullptr) (void)hipEventDestroy(event);
   free_signal(impl_->ready);
   free_signal(impl_->done);
 }
@@ -244,6 +247,13 @@ bool GpuSync::probe(uintptr_t /*stream*/) {
 void GpuSync::gpu_submit(uintptr_t stream, uint64_t task) {
   auto s = reinterpret_cast<hipStream_t>(stream);
   hipStreamBatchMemOpParams* ops = needs_explicit_nodes(s) ? impl_->submit_ops[task].data() : nullptr;
+  // Keep each task's event alive across graph capture, replay and recapture.
+  auto& publication = impl_->publication_events[task];
+  if (publication == nullptr)
+    hip_check(hipEventCreateWithFlags(&publication, hipEventDisableTiming | hipEventReleaseToSystem),
+              "create input publication event");
+  // External recording retains the D2H publication point in every graph replay.
+  hip_check(hipEventRecordWithFlags(publication, s, hipEventRecordExternal), "publish input copies");
   // Reset done before raising ready, so the coordinator's completion for this step can never be wiped.
   hip_check(write_value(s, impl_->done, 0, ops), "write done");
   hip_check(write_value(s, impl_->ready, task, ops ? ops + 1 : nullptr), "write ready");

python/freetoken/moe/cpu_executor.py

@@ -5,11 +5,12 @@
 the whole CUDA-graph-capturable sequence on the current stream:
 
     D2H (hidden, topk_ids, topk_weights -> pinned IO rows [:bs])
+      -> publish (HIP external event before the memop doorbell)
       -> submit (doorbell naming (layer, bs), or a cudaLaunchHostFunc host node)
       -> sync   (wait for the CPU pool to finish)
       -> H2D (pinned expert output -> GPU)
 
-With the doorbell on CUDA, a publish kernel does the D2H and the doorbell, and skips the step when no route goes to the CPU; a combine kernel does the H2D (csrc/cpu_moe/gpu_sync_io.h).
+With kernel IO, a publish kernel does the D2H and the doorbell, and skips the step when no route goes to the CPU; a combine kernel does the H2D (csrc/cpu_moe/gpu_sync_io.h).
 """
 
 from __future__ import annotations
@@ -171,6 +172,8 @@ def __init__(
         num_threads: int,
         max_tokens: int,
         device: torch.device,
+        activation_dtype: torch.dtype = torch.bfloat16,
+        graph_enabled: bool = True,
         swiglu_alpha: float = 1.702,
         swiglu_limit: float | None = None,
         layer_formats: list[tuple[str, dict[str, int]]] | None = None,
@@ -342,14 +345,12 @@ def __init__(
         self._kernel_io = False
         if self._flag_sync:
             self._ext.start_flag_coordinator(self._coord_core)
-            # ROCm keeps the copies until the kernel IO has been checked on AMD GPUs.
-            if torch.version.hip is None:
-                reason = _cpu_moe.kernel_io_probe(torch.cuda.current_stream().cuda_stream)
-                if reason:
-                    logger.warning_rank0(f"cpu-moe kernel IO unavailable ({reason}); using D2H/H2D copies around the doorbell")
-                else:
-                    self._ext.enable_kernel_io()
-                    self._kernel_io = True
+            reason = _cpu_moe.kernel_io_probe(torch.cuda.current_stream().cuda_stream)
+            if reason:
+                logger.warning_rank0(f"cpu-moe kernel IO unavailable ({reason}); using D2H/H2D copies around the doorbell")
+            else:
+                self._ext.enable_kernel_io()
+                self._kernel_io = True
             self._watchdog_stop = False
             # The thread target holds a WEAKREF and re-derefs it each tick: a bound
             # method would strong-reference the executor forever (the loop never ends
@@ -366,6 +367,21 @@ def __init__(
             )
             self._watchdog.start()
 
+        self._graph_safe_dtypes = {activation_dtype}
+        if torch.version.hip is not None:
+            self._graph_capture_safe = self._flag_sync
+            if self._flag_sync:
+                reason = self._probe_graph_io(activation_dtype, capture_graph=graph_enabled)
+                if reason and self._kernel_io and activation_dtype == torch.bfloat16:
+                    logger.warning_rank0(f"cpu-moe kernel IO probe failed ({reason}); checking event + memcpy IO")
+                    # Python selects the copy path; the native KernelIo allocation lives until executor teardown.
+                    self._kernel_io = False
+                    reason = self._probe_graph_io(activation_dtype, capture_graph=graph_enabled)
+                if reason:
+                    if not self._eager_io_safe:
+                        raise RuntimeError(f"CPU MoE eager IO probe failed for {activation_dtype}: {reason}")
+                    logger.warning_rank0(f"cpu-moe graph IO probe failed for {activation_dtype} ({reason}); using eager decode")
+                self._graph_capture_safe = graph_enabled and not reason
 
         fmt_label = "; ".join(f"{_format_label(name, opts)}: {n} layers" for (name, opts), n in layer_counts.items())
         logger.info_rank0(
@@ -375,9 +391,88 @@ def __init__(
             f"H={self.H} I={self.I} experts={self.num_experts} layers={self.num_layers} "
             f"top_k={self.top_k} act={activation} max_tokens={self.max_tokens} "
             f"spin_wait={self.spin_wait} "
-            f"handshake={'kernel' if self._kernel_io else 'memop' if self._flag_sync else 'host-func'}"
+            f"handshake={'kernel' if self._kernel_io else 'event-memop' if self._flag_sync and torch.version.hip is not None else 'memop' if self._flag_sync else 'host-func'}"
         )
 
+    def graph_capture_safe(self, dtype: torch.dtype) -> bool:
+        return self._graph_capture_safe and (
+            torch.version.hip is None or dtype in self._graph_safe_dtypes
+        )
+
+    def _probe_graph_io(self, dtype: torch.dtype, *, capture_graph: bool = True) -> str:
+        """Qualify IO against synchronous CPU execution, including hybrid accumulation."""
+        self._eager_io_safe = False
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            bs = min(2, self.max_tokens)
+            hidden = torch.empty(3, bs, self.H, device="cuda", dtype=dtype)
+            ids = torch.empty(3, bs, self.top_k, device="cuda", dtype=torch.int32)
+            weights = torch.empty(3, bs, self.top_k, device="cuda", dtype=torch.float32)
+            add = torch.empty_like(hidden)
+            base = torch.arange(hidden.numel(), dtype=torch.float32).reshape(hidden.shape)
+            routes = torch.arange(ids.numel()).reshape(ids.shape)
+
+            def inputs(step):
+                hidden.copy_(((base + step * 7) % 31 - 15) / 64)
+                cpu_ids = ((routes + step) % self.num_experts).int()
+                cpu_ids[step % 3].fill_(-1)
+                cpu_ids[(step + 1) % 3, :, -1] = -1
+                ids.copy_(cpu_ids)
+                weights.copy_(((routes + step) % 7 + 1).float() / (7 * self.top_k))
+                add.copy_(((base + step * 3) % 17 - 8) / 32)
+
+            def forward():
+                outputs = []
+                for i in range(3):
+                    pending = self.decode_submit(i % self.num_layers, hidden[i], weights[i], ids[i])
+                    outputs.append(self.decode_sync(pending, add_to=add[i].clone() if i == 1 else None))
+                return outputs
+
+            def reference():
+                # No asynchronous IO or notification: both paths can otherwise share the same bug.
+                x = hidden
+                if self._gpu_prequant:
+                    from freetoken.kernel.triton.dsv4.fp8_linear import act_quant_fp8_roundtrip
+
+                    x = torch.stack([act_quant_fp8_roundtrip(row, block=self._act_block) for row in hidden])
+                host_x, host_ids, host_w, host_add = x.cpu(), ids.cpu(), weights.cpu(), add.cpu()
+                expected = []
+                for i in range(3):
+                    self._io["x"][:bs].copy_(host_x[i])
+                    self._io["ids"][:bs].copy_(host_ids[i])
+                    self._io["w"][:bs].copy_(host_w[i])
+                    self._ext.run(i % self.num_layers, bs)
+                    # Match the IO rounding boundary before hybrid addition; the BF16 combine kernel
+                    # also adds already-rounded BF16 operands in FP32, then rounds back to BF16.
+                    out = self._io["y"][:bs].to(dtype).clone()
+                    expected.append(out + host_add[i] if i == 1 else out)
+                return expected
+
+            def matches(actual, expected):
+                return all(torch.equal(got.cpu(), want) for got, want in zip(actual, expected))
+
+            inputs(0)
+            eager = forward()
+            stream.synchronize()
+            if not matches(eager, reference()):
+                return "eager output differs from synchronous CPU reference"
+            self._eager_io_safe = True
+            if not capture_graph:
+                return ""
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, stream=stream):
+                captured = forward()
+            for step in range(1, 4):
+                inputs(step)
+                graph.replay()
+                stream.synchronize()
+                self.raise_if_unhealthy()
+                # Read the reference only after replay, so it cannot prefill stale graph IO.
+                if not matches(captured, reference()):
+                    return "graph output differs from synchronous CPU reference"
+        return ""
+
     def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor:
         """Build a CPU int64 tensor of per-layer base addresses for one bank.
 
@@ -559,11 +654,10 @@ def decode_submit(
         bs = hidden_states.shape[0]
         if bs > self.max_tokens:
             raise ValueError(f"CPU MoE decode batch {bs} exceeds the executor's max_tokens {self.max_tokens}")
-        if not self._graph_capture_safe and torch.cuda.is_current_stream_capturing():
+        if not self.graph_capture_safe(hidden_states.dtype) and torch.cuda.is_current_stream_capturing():
             raise RuntimeError(
-                "CPU MoE decode cannot be captured into a CUDA graph here: the stream-memop "
-                "doorbell is off or failed its startup probe, and the host-function sync is not "
-                "replay-safe on this runtime; run with --cuda-graph-max-bs 0"
+                "CPU MoE decode cannot be captured into a CUDA graph here: the IO/handshake "
+                "path is not replay-safe for this runtime and dtype; run with --cuda-graph-max-bs 0"
             )
 
         if self._gpu_prequant:
@@ -592,7 +686,7 @@ def decode_submit(
         out = torch.empty_like(hidden_states)
         if self._flag_sync:
             # Front-end memops: reset done, then ring the doorbell with (layer, bs).
-            # No kernel launched; no host-func round trip.
+            # HIP records the input-publication event before these writes.
             self._ext.flag_signal(stream, layer_id, bs)
         else:
             self._ext.submit_with_cuda_stream(stream, layer_id, bs)

python/freetoken/moe/locked_fetch.py

@@ -19,8 +19,8 @@
 
 
 def locked_fetch_wanted(config) -> bool:
-    """Whether locked layers fetch through pinned staging: hybrid, --moe-locked-fetch not 0, CUDA (the copy kernels are untested on ROCm)."""
-    return config.moe_strategy == "hybrid" and config.moe_locked_fetch != 0 and torch.version.hip is None
+    """Whether hybrid layers fetch through pinned staging instead of staying CPU-only."""
+    return config.moe_strategy == "hybrid" and config.moe_locked_fetch != 0
 
 
 def check_locked_fetch_config(config) -> None:

tests/engine/test_moe_cpu_layers.py

@@ -61,6 +61,27 @@ def test_resolve_backend_dispatch():
     assert resolve(_cfg("fused", "8"), L) == frozenset()
 
 
+@pytest.mark.parametrize("local_safe,remote_safe", [(True, True), (False, True), (True, False)])
+def test_cpu_moe_graph_fallback_agrees_across_tp_ranks(monkeypatch, local_safe, remote_safe):
+    import torch
+    from freetoken.engine.engine import _disable_unsafe_cpu_moe_graphs
+
+    config = SimpleNamespace(tp_info=SimpleNamespace(size=2), cuda_graph_bs=[1, 4], cuda_graph_max_bs=4)
+    executor = SimpleNamespace(graph_capture_safe=lambda dtype: local_safe)
+    group = object()
+    calls = []
+
+    def reduce(tensor, *, op, group):
+        calls.append((op, group))
+        tensor.fill_(bool(tensor.item()) and remote_safe)
+
+    monkeypatch.setattr(torch.distributed, "all_reduce", reduce)
+    _disable_unsafe_cpu_moe_graphs(config, executor, torch.bfloat16, group)
+    assert calls == [(torch.distributed.ReduceOp.MIN, group)]
+    assert config.cuda_graph_bs == ([1, 4] if local_safe and remote_safe else [])
+    assert config.cuda_graph_max_bs == (4 if local_safe and remote_safe else 0)
+
+
 if __name__ == "__main__":
     import sys
 

tests/moe/hip_probe_faults.cpp

@@ -34,6 +34,14 @@ extern "C" hipError_t hipStreamWriteValue64(hipStream_t stream, void* address, u
   return real<decltype(&hipStreamWriteValue64)>("hipStreamWriteValue64")(stream, address, value, flags);
 }
 
+extern "C" hipError_t hipEventRecordWithFlags(hipEvent_t event, hipStream_t stream, unsigned flags) {
+  if (mode("publication-event") && flags == hipEventRecordExternal) {
+    ++injected;
+    return hipErrorNotSupported;
+  }
+  return real<decltype(&hipEventRecordWithFlags)>("hipEventRecordWithFlags")(event, stream, flags);
+}
+
 extern "C" hipError_t hipStreamUpdateCaptureDependencies(
     hipStream_t stream, hipGraphNode_t* dependencies, size_t count, unsigned flags) {
   if (mode("dependency")) {