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

docs/install_amd.md

@@ -49,7 +49,18 @@ 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. Before notifying the CPU,
+FreeToken records a system-release event after copying the expert inputs to
+host memory. The event is recorded as an external graph node so that every
+replay retains the publication point. This application workaround keeps Graph
+and memops enabled on unpatched runtimes; it does not require a custom HIP library.
+
+Native memop capture is detected at runtime. If it is unavailable, FreeToken
+adds explicit memop graph nodes. Before model capture, a numerical check
+compares distinct tasks and changing inputs against eager execution. ROCm
+decode graphs require BF16 inputs and a passing handshake and data check. When
+these checks reject the path, the engine logs a warning and disables decode
+graphs across TP ranks. Capture and runtime API errors still stop startup.
+
+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

@@ -40,6 +40,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
@@ -963,6 +974,7 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None:
         )
         cache.set_cpu_executor(executor)
         self.cpu_moe_executor = executor
+        _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,6 +157,7 @@ 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{};
@@ -175,6 +176,8 @@ GpuSync::GpuSync() : impl_(std::make_unique<Impl>()) {
 }
 
 GpuSync::~GpuSync() {
+  for (const auto& [task, event] : impl_->publication_events)
+    if (event != nullptr) (void)hipEventDestroy(event);
   free_signal(impl_->ready);
   free_signal(impl_->done);
 }
@@ -243,6 +246,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
@@ -342,7 +343,6 @@ 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:
@@ -366,6 +366,14 @@ def __init__(
             )
             self._watchdog.start()
 
+        if torch.version.hip is not None:
+            # Check the event + memcpy + memop protocol, beyond the flag-only probe.
+            self._graph_capture_safe = self._flag_sync
+            if self._graph_capture_safe:
+                reason = self._probe_graph_io()
+                if reason:
+                    self._graph_capture_safe = False
+                    logger.warning_rank0(f"cpu-moe graph IO probe failed ({reason}); using eager decode")
 
         fmt_label = "; ".join(f"{_format_label(name, opts)}: {n} layers" for (name, opts), n in layer_counts.items())
         logger.info_rank0(
@@ -375,9 +383,54 @@ 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:
+        # HIP graph IO is qualified for BF16 inputs.
+        return self._graph_capture_safe and (
+            torch.version.hip is None or dtype == torch.bfloat16
         )
 
+    def _probe_graph_io(self) -> str:
+        """Check distinct tasks sharing IO, including CPU-route skips, before engine capture."""
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            hidden = torch.empty(3, 1, self.H, device="cuda", dtype=torch.bfloat16)
+            ids = torch.empty(3, 1, self.top_k, device="cuda", dtype=torch.int32)
+            weights = torch.empty(3, 1, self.top_k, device="cuda", dtype=torch.float32)
+            base = torch.arange(3 * self.H, dtype=torch.float32).reshape(3, 1, self.H)
+            routes = torch.arange(3 * self.top_k).reshape(3, 1, self.top_k)
+
+            def inputs(step):
+                hidden.copy_((((base + step * 7) % 31 - 15) / 64).bfloat16())
+                cpu_ids = ((routes + step) % self.num_experts).int()
+                cpu_ids[step % 3].fill_(-1)
+                ids.copy_(cpu_ids)
+                weights.copy_(((routes + step) % 7 + 1).float() / (7 * self.top_k))
+
+            def forward():
+                return [self.decode(i % self.num_layers, hidden[i], weights[i], ids[i]) for i in range(3)]
+
+            inputs(0)
+            forward()
+            stream.synchronize()
+            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()
+                # Eager must follow replay: otherwise it can prefill shared IO and hide stale graph inputs.
+                eager = forward()
+                stream.synchronize()
+                self.raise_if_unhealthy()
+                if any(not torch.equal(got, expected) for got, expected in zip(captured, eager)):
+                    return "graph output differs from eager with distinct task inputs"
+        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 +612,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 +644,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)

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")) {

tests/moe/test_cpu_moe.py

@@ -545,6 +545,93 @@ def test_cpu_moe_hip_probe_error_recovery(hip_probe_fault_library, fault):
     assert result.returncode == 0, result.stdout + result.stderr
 
 
+@pytest.mark.skipif(torch.version.hip is None, reason="requires ROCm")
+@pytest.mark.parametrize("fallback", ["handshake-probe", "graph-data", "fp16"])
+def test_cpu_moe_hip_unsafe_graph_keeps_eager_usable(monkeypatch, fallback):
+    from freetoken.kernel import _cpu_moe
+    from freetoken.moe.cpu_executor import CpuMoeExecutor
+
+    if fallback == "handshake-probe":
+        monkeypatch.setattr(_cpu_moe, "memops_probe", lambda stream: False)
+    elif fallback == "graph-data":
+        original = CpuMoeExecutor.decode_sync
+
+        def corrupt_capture(self, *args, **kwargs):
+            out = original(self, *args, **kwargs)
+            return out + 1 if torch.cuda.is_current_stream_capturing() else out
+
+        monkeypatch.setattr(CpuMoeExecutor, "decode_sync", corrupt_capture)
+
+    cache = _make_cache(3, 8, 256, 128, scale=0.03)
+    ex = CpuMoeExecutor(
+        cache, top_k=4, activation="silu", apply_router_weight_on_input=False,
+        num_threads=4, max_tokens=1, device=torch.device("cuda"),
+    )
+    dtype = torch.float16 if fallback == "fp16" else torch.bfloat16
+    assert not ex.graph_capture_safe(dtype)
+    if fallback == "graph-data":
+        assert ex._flag_sync and not ex._kernel_io
+    hidden = torch.randn(1, 256, device="cuda", dtype=torch.bfloat16).to(dtype)
+    ids = torch.tensor([[0, 2, -1, 5]], device="cuda", dtype=torch.int32)
+    weights = torch.tensor([[0.2, 0.3, 0.1, 0.4]], device="cuda")
+    actual = ex.decode(1, hidden, weights, ids)
+    expected = _reference(cache, 1, hidden, ids.cpu(), weights.cpu())
+    assert actual.dtype == dtype
+    relative = (actual.float().cpu() - expected).abs().max() / (expected.abs().max() + 1e-6)
+    assert relative < 2e-2, relative
+    with pytest.raises(RuntimeError, match="--cuda-graph-max-bs 0"):
+        with torch.cuda.graph(torch.cuda.CUDAGraph(), stream=torch.cuda.Stream()):
+            ex.decode(1, hidden, weights, ids)
+    assert torch.isfinite(ex.decode(1, hidden, weights, ids).cpu()).all()
+
+
+@pytest.mark.skipif(torch.version.hip is None, reason="requires ROCm")
+def test_cpu_moe_hip_event_path_does_not_require_kernel_io(monkeypatch):
+    from freetoken.kernel import _cpu_moe
+    from freetoken.moe.cpu_executor import CpuMoeExecutor
+
+    def reject_kernel_probe(stream):
+        raise AssertionError("HIP event publication must not depend on the publish kernel")
+
+    monkeypatch.setattr(_cpu_moe, "kernel_io_probe", reject_kernel_probe)
+    cache = _make_cache(3, 8, 256, 128, scale=0.03)
+    ex = CpuMoeExecutor(
+        cache, top_k=4, activation="silu", apply_router_weight_on_input=False,
+        num_threads=4, max_tokens=1, device=torch.device("cuda"),
+    )
+    assert ex._flag_sync and not ex._kernel_io
+    assert ex.graph_capture_safe(torch.bfloat16)
+
+
+def test_cpu_moe_hip_event_record_failure_is_loud(hip_probe_fault_library):
+    import os
+    from pathlib import Path
+    import subprocess
+    import sys
+
+    script = r"""
+import ctypes, runpy, sys
+import pytest
+import torch
+from freetoken.moe.cpu_executor import CpuMoeExecutor
+helpers = runpy.run_path(sys.argv[1])
+shim = ctypes.CDLL(sys.argv[2])
+cache = helpers['_make_cache'](3, 8, 256, 128, scale=0.03)
+with pytest.raises(RuntimeError, match='publish input copies'):
+    CpuMoeExecutor(cache, top_k=4, activation='silu', apply_router_weight_on_input=False,
+        num_threads=4, max_tokens=1, device=torch.device('cuda'))
+torch.cuda.synchronize()
+assert shim.probe_test_injected() > 0
+"""
+    env = dict(os.environ, LD_PRELOAD=str(hip_probe_fault_library),
+               FREETOKEN_TEST_HIP_PROBE_FAULT="publication-event", FREETOKEN_CPU_MOE_FLAG_SYNC="1")
+    result = subprocess.run(
+        [sys.executable, "-c", script, str(Path(__file__).resolve()), str(hip_probe_fault_library)],
+        env=env, capture_output=True, text=True, timeout=30,
+    )
+    assert result.returncode == 0, result.stdout + result.stderr
+
+
 def test_cpu_moe_decode_cuda_graph_replay():
     from freetoken.moe.cpu_executor import CpuMoeExecutor