From 616f68d9c7ae71d176febbfa83f1e8bfa0468746 Mon Sep 17 00:00:00 2001
From: zihaomu <zihaomu@outlook.com>
Date: Tue, 20 Oct 2026 20:11:03 +0800
Subject: [PATCH 1/4] fix(moe): prevent stale cpu expert inputs in rocm graphs

---
 python/freetoken/engine/engine.py    | 12 +++++
 python/freetoken/moe/cpu_executor.py | 76 +++++++++++++++++++++++-----
 tests/engine/test_moe_cpu_layers.py  | 21 ++++++++
 tests/moe/test_cpu_moe.py            | 40 +++++++++++++++
 tests/moe/test_cpu_moe_fp8_block.py  | 52 +++++++++++++++++++
 5 files changed, 188 insertions(+), 13 deletions(-)

diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 590cf80cc..f4a1901b0 100644
--- a/python/freetoken/engine/engine.py
+++ b/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
@@ -1093,6 +1104,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."""
diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py
index 1f275749d..8f4585516 100644
--- a/python/freetoken/moe/cpu_executor.py
+++ b/python/freetoken/moe/cpu_executor.py
@@ -9,7 +9,7 @@
       -> 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,14 +342,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 +364,14 @@ def __init__(
             )
             self._watchdog.start()
 
+        if torch.version.hip is not None:
+            # The flag probe alone misses stale shared IO in HIP graphs using memcpy + memops.
+            self._graph_capture_safe = self._flag_sync and self._kernel_io
+            if self._graph_capture_safe:
+                reason = self._probe_kernel_graph_io()
+                if reason:
+                    self._graph_capture_safe = False
+                    logger.warning_rank0(f"cpu-moe kernel IO graph 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(
@@ -378,6 +384,51 @@ def __init__(
             f"handshake={'kernel' if self._kernel_io else 'memop' if self._flag_sync else 'host-func'}"
         )
 
+    def graph_capture_safe(self, dtype: torch.dtype) -> bool:
+        # Non-bf16 inputs still take the memcpy path, even when kernel IO is enabled.
+        return self._graph_capture_safe and (
+            torch.version.hip is None or (self._kernel_io and dtype == torch.bfloat16)
+        )
+
+    def _probe_kernel_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 +610,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:
diff --git a/tests/engine/test_moe_cpu_layers.py b/tests/engine/test_moe_cpu_layers.py
index feb19300b..2bfee2770 100644
--- a/tests/engine/test_moe_cpu_layers.py
+++ b/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
 
diff --git a/tests/moe/test_cpu_moe.py b/tests/moe/test_cpu_moe.py
index edd9ddc9d..04d3b6f2d 100644
--- a/tests/moe/test_cpu_moe.py
+++ b/tests/moe/test_cpu_moe.py
@@ -545,6 +545,46 @@ 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", ["kernel-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 == "kernel-probe":
+        monkeypatch.setattr(_cpu_moe, "kernel_io_probe", lambda stream: "injected kernel IO rejection")
+    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._kernel_io, "must test rejection of the full data protocol, not the kernel capability probe"
+    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()
+
+
 def test_cpu_moe_decode_cuda_graph_replay():
     from freetoken.moe.cpu_executor import CpuMoeExecutor
 
diff --git a/tests/moe/test_cpu_moe_fp8_block.py b/tests/moe/test_cpu_moe_fp8_block.py
index baecb86f9..f6cd4ef7b 100644
--- a/tests/moe/test_cpu_moe_fp8_block.py
+++ b/tests/moe/test_cpu_moe_fp8_block.py
@@ -261,3 +261,55 @@ def test_fp16_bits_carry_every_e4m3_byte():
     finite = torch.isfinite(ref)
     assert torch.equal(got[finite], ref[finite])
     assert torch.equal(got[~finite].abs(), torch.full((2,), 480.0))  # 0x7F / 0xFF, as e4m3_decode
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+@pytest.mark.parametrize("bs", [1, 3])
+@pytest.mark.parametrize("accumulate", [False, True])
+def test_cpu_fp8_block_graph_distinct_layer_inputs(bs, accumulate):
+    from freetoken.moe.cpu_executor import CpuMoeExecutor
+
+    L, E, H, I, top_k = 3, 8, 256, 128, 4
+    cache = _make_fp8_block_cache(L, E, H, I, seed=7)
+    stream = torch.cuda.Stream()
+    with torch.cuda.stream(stream):
+        ex = CpuMoeExecutor(
+            cache, top_k=top_k, activation="silu", apply_router_weight_on_input=False,
+            num_threads=4, max_tokens=bs, device=torch.device("cuda"),
+        )
+        hidden = torch.randn(L, bs, H, device="cuda", dtype=torch.bfloat16)
+        ids = torch.zeros(L, bs, top_k, device="cuda", dtype=torch.int32)
+        weights = torch.rand(L, bs, top_k, device="cuda")
+        gpu_out = torch.randn_like(hidden)
+
+        def forward():
+            outputs = []
+            for layer in range(L):
+                pending = ex.decode_submit(layer, hidden[layer], weights[layer], ids[layer])
+                outputs.append(ex.decode_sync(pending, add_to=gpu_out[layer].clone() if accumulate else None))
+            return outputs
+
+        forward()
+        stream.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            captured = forward()
+
+        for step in range(3):
+            torch.manual_seed(101 + step)
+            hidden.copy_(torch.randn(L, bs, H, dtype=torch.bfloat16))
+            new_ids = torch.randint(0, E, (L, bs, top_k), dtype=torch.int32)
+            new_ids[step].fill_(-1)
+            new_ids[(step + 1) % L, :, 0] = -1
+            ids.copy_(new_ids)
+            weights.copy_(torch.rand(L, bs, top_k))
+            gpu_out.copy_(torch.randn(L, bs, H, dtype=torch.bfloat16))
+            # Do not run eager first: it would prefill the shared host IO with this step's data.
+            graph.replay()
+            stream.synchronize()
+            eager = forward()
+            stream.synchronize()
+            for got, expected in zip(captured, eager):
+                torch.testing.assert_close(got, expected, rtol=0, atol=0)
+            skipped = gpu_out[step] if accumulate else torch.zeros_like(gpu_out[step])
+            torch.testing.assert_close(captured[step], skipped, rtol=0, atol=0)

From b51c1a87367a321d87ca622aeb4246176b4cb98a Mon Sep 17 00:00:00 2001
From: zihaomu <zihaomu@outlook.com>
Date: Tue, 27 Oct 2026 19:29:09 +0800
Subject: [PATCH 2/4] fix(moe): publish rocm cpu inputs with graph events

---
 docs/install_amd.md                           | 19 +++++--
 .../kernel/csrc/cpu_moe/gpu_sync_hip.cpp      | 10 ++++
 python/freetoken/moe/cpu_executor.py          | 32 ++++++-----
 tests/moe/hip_probe_faults.cpp                |  8 +++
 tests/moe/test_cpu_moe.py                     | 55 +++++++++++++++++--
 tests/moe/test_cpu_moe_fp8_block.py           | 49 +++++++++++++++++
 6 files changed, 150 insertions(+), 23 deletions(-)

diff --git a/docs/install_amd.md b/docs/install_amd.md
index ed934ad65..6bfca41cb 100644
--- a/docs/install_amd.md
+++ b/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.
diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp
index 395f26841..c834f4044 100644
--- a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp
+++ b/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{};
@@ -176,6 +177,8 @@ GpuSync::GpuSync(bool) : 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);
 }
@@ -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");
diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py
index 8f4585516..2fe908bb6 100644
--- a/python/freetoken/moe/cpu_executor.py
+++ b/python/freetoken/moe/cpu_executor.py
@@ -5,6 +5,7 @@
 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)
@@ -342,12 +343,13 @@ def __init__(
         self._kernel_io = False
         if self._flag_sync:
             self._ext.start_flag_coordinator(self._coord_core)
-            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
+            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
             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
@@ -365,13 +367,13 @@ def __init__(
             self._watchdog.start()
 
         if torch.version.hip is not None:
-            # The flag probe alone misses stale shared IO in HIP graphs using memcpy + memops.
-            self._graph_capture_safe = self._flag_sync and self._kernel_io
+            # 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_kernel_graph_io()
+                reason = self._probe_graph_io()
                 if reason:
                     self._graph_capture_safe = False
-                    logger.warning_rank0(f"cpu-moe kernel IO graph probe failed ({reason}); using eager decode")
+                    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(
@@ -381,16 +383,16 @@ 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:
-        # Non-bf16 inputs still take the memcpy path, even when kernel IO is enabled.
+        # HIP graph IO is qualified for BF16 inputs.
         return self._graph_capture_safe and (
-            torch.version.hip is None or (self._kernel_io and dtype == torch.bfloat16)
+            torch.version.hip is None or dtype == torch.bfloat16
         )
 
-    def _probe_kernel_graph_io(self) -> str:
+    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())
@@ -642,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)
diff --git a/tests/moe/hip_probe_faults.cpp b/tests/moe/hip_probe_faults.cpp
index 4426bf0d1..fe3dfbae1 100644
--- a/tests/moe/hip_probe_faults.cpp
+++ b/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")) {
diff --git a/tests/moe/test_cpu_moe.py b/tests/moe/test_cpu_moe.py
index 04d3b6f2d..26ed89c6f 100644
--- a/tests/moe/test_cpu_moe.py
+++ b/tests/moe/test_cpu_moe.py
@@ -546,13 +546,13 @@ def test_cpu_moe_hip_probe_error_recovery(hip_probe_fault_library, fault):
 
 
 @pytest.mark.skipif(torch.version.hip is None, reason="requires ROCm")
-@pytest.mark.parametrize("fallback", ["kernel-probe", "graph-data", "fp16"])
+@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 == "kernel-probe":
-        monkeypatch.setattr(_cpu_moe, "kernel_io_probe", lambda stream: "injected kernel IO rejection")
+    if fallback == "handshake-probe":
+        monkeypatch.setattr(_cpu_moe, "memops_probe", lambda stream: False)
     elif fallback == "graph-data":
         original = CpuMoeExecutor.decode_sync
 
@@ -570,7 +570,7 @@ def corrupt_capture(self, *args, **kwargs):
     dtype = torch.float16 if fallback == "fp16" else torch.bfloat16
     assert not ex.graph_capture_safe(dtype)
     if fallback == "graph-data":
-        assert ex._kernel_io, "must test rejection of the full data protocol, not the kernel capability probe"
+        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")
@@ -585,6 +585,53 @@ def corrupt_capture(self, *args, **kwargs):
     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
 
diff --git a/tests/moe/test_cpu_moe_fp8_block.py b/tests/moe/test_cpu_moe_fp8_block.py
index f6cd4ef7b..8f25b5ee8 100644
--- a/tests/moe/test_cpu_moe_fp8_block.py
+++ b/tests/moe/test_cpu_moe_fp8_block.py
@@ -313,3 +313,52 @@ def forward():
                 torch.testing.assert_close(got, expected, rtol=0, atol=0)
             skipped = gpu_out[step] if accumulate else torch.zeros_like(gpu_out[step])
             torch.testing.assert_close(captured[step], skipped, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(torch.version.hip is None, reason="requires ROCm")
+def test_cpu_fp8_block_event_graph_replay_and_recapture():
+    from freetoken.moe.cpu_executor import CpuMoeExecutor
+
+    L, E, H, I, top_k = 3, 8, 256, 128, 4
+    cache = _make_fp8_block_cache(L, E, H, I, seed=71)
+    ex = CpuMoeExecutor(
+        cache, top_k=top_k, activation="silu", apply_router_weight_on_input=False,
+        num_threads=4, max_tokens=3, device=torch.device("cuda"),
+    )
+    assert ex._flag_sync and not ex._kernel_io
+    assert ex.graph_capture_safe(torch.bfloat16)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        # Reuse the executor after destroying each graph, including a repeated task size.
+        for bs in (1, 3, 1):
+            hidden = torch.empty(L, bs, H, device="cuda", dtype=torch.bfloat16)
+            ids = torch.empty(L, bs, top_k, device="cuda", dtype=torch.int32)
+            weights = torch.empty(L, bs, top_k, device="cuda", dtype=torch.float32)
+
+            def inputs(step):
+                generator = torch.Generator().manual_seed(1700 + step)
+                hidden.copy_(torch.randn(hidden.shape, dtype=hidden.dtype, generator=generator))
+                routes = torch.randint(E, ids.shape, dtype=ids.dtype, generator=generator)
+                routes[step % L].fill_(-1)
+                ids.copy_(routes)
+                weights.copy_(torch.rand(weights.shape, generator=generator))
+
+            def forward():
+                return [ex.decode(layer, hidden[layer], weights[layer], ids[layer]) for layer in range(L)]
+
+            inputs(0)
+            # New tasks may allocate their publication event inside capture.
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, stream=stream):
+                outputs = forward()
+            for step in range(1, 11):
+                inputs(step)
+                graph.replay()
+                stream.synchronize()
+                references = forward()
+                stream.synchronize()
+                for actual, expected in zip(outputs, references):
+                    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+            graph.reset()
+    ex.raise_if_unhealthy()

From fc141047055caf7f87117b6a1403b53cc54934d8 Mon Sep 17 00:00:00 2001
From: LZ-QWQ <lz1830837889@163.com>
Date: Fri, 30 Oct 2026 18:17:43 +0800
Subject: [PATCH 3/4] fix(moe): qualify rocm kernel io with ordered copy
 fallback

---
 docs/install_amd.md                  |  31 +++--
 python/freetoken/engine/engine.py    |  10 +-
 python/freetoken/moe/cpu_executor.py | 100 +++++++++++-----
 tests/moe/test_cpu_moe.py            | 168 +++++++++++++++++++++++++--
 tests/moe/test_cpu_moe_fp8_block.py  |  10 +-
 5 files changed, 265 insertions(+), 54 deletions(-)

diff --git a/docs/install_amd.md b/docs/install_amd.md
index 6bfca41cb..fed29d0ec 100644
--- a/docs/install_amd.md
+++ b/docs/install_amd.md
@@ -49,18 +49,31 @@ 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. 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.
+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 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.
+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.
 
 After updating the native CPU MoE sources, rebuild the extension with
 `python setup.py build_ext --inplace` in the ROCm environment before serving.
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index f4a1901b0..7c9e5102c 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -1088,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(
@@ -1098,13 +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
-        _disable_unsafe_cpu_moe_graphs(config, executor, self.dtype, self.tp_cpu_group)
+        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."""
diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py
index 2fe908bb6..fea24b40b 100644
--- a/python/freetoken/moe/cpu_executor.py
+++ b/python/freetoken/moe/cpu_executor.py
@@ -172,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,
@@ -343,13 +345,12 @@ def __init__(
         self._kernel_io = False
         if self._flag_sync:
             self._ext.start_flag_coordinator(self._coord_core)
-            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,14 +367,21 @@ def __init__(
             )
             self._watchdog.start()
 
+        self._graph_safe_dtypes = {activation_dtype}
         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 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:
-                    self._graph_capture_safe = False
-                    logger.warning_rank0(f"cpu-moe graph IO probe failed ({reason}); using eager decode")
+                    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(
@@ -387,35 +395,71 @@ def __init__(
         )
 
     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
+            torch.version.hip is None or dtype in self._graph_safe_dtypes
         )
 
-    def _probe_graph_io(self) -> str:
-        """Check distinct tasks sharing IO, including CPU-route skips, before engine capture."""
+    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):
-            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)
+            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).bfloat16())
+                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():
-                return [self.decode(i % self.num_layers, hidden[i], weights[i], ids[i]) for i in range(3)]
+                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)
-            forward()
+            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()
@@ -423,12 +467,10 @@ def forward():
                 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"
+                # 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:
diff --git a/tests/moe/test_cpu_moe.py b/tests/moe/test_cpu_moe.py
index 26ed89c6f..ed87a3cc7 100644
--- a/tests/moe/test_cpu_moe.py
+++ b/tests/moe/test_cpu_moe.py
@@ -546,7 +546,7 @@ def test_cpu_moe_hip_probe_error_recovery(hip_probe_fault_library, fault):
 
 
 @pytest.mark.skipif(torch.version.hip is None, reason="requires ROCm")
-@pytest.mark.parametrize("fallback", ["handshake-probe", "graph-data", "fp16"])
+@pytest.mark.parametrize("fallback", ["handshake-probe", "graph-data"])
 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
@@ -567,7 +567,7 @@ def corrupt_capture(self, *args, **kwargs):
         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
+    dtype = torch.bfloat16
     assert not ex.graph_capture_safe(dtype)
     if fallback == "graph-data":
         assert ex._flag_sync and not ex._kernel_io
@@ -586,14 +586,11 @@ def corrupt_capture(self, *args, **kwargs):
 
 
 @pytest.mark.skipif(torch.version.hip is None, reason="requires ROCm")
-def test_cpu_moe_hip_event_path_does_not_require_kernel_io(monkeypatch):
+def test_cpu_moe_hip_event_fallback_when_kernel_io_unavailable(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)
+    monkeypatch.setattr(_cpu_moe, "kernel_io_probe", lambda stream: "mapped IO unavailable")
     cache = _make_cache(3, 8, 256, 128, scale=0.03)
     ex = CpuMoeExecutor(
         cache, top_k=4, activation="silu", apply_router_weight_on_input=False,
@@ -603,30 +600,177 @@ def reject_kernel_probe(stream):
     assert ex.graph_capture_safe(torch.bfloat16)
 
 
-def test_cpu_moe_hip_event_record_failure_is_loud(hip_probe_fault_library):
+@pytest.mark.skipif(torch.version.hip is None, reason="requires ROCm")
+@pytest.mark.parametrize("graph_bs,max_bs", [(None, 0), ([], 4), (None, 1), ([1], 0), (None, None), (None, -1)])
+def test_cpu_moe_hip_kernel_graph_mismatch_uses_event(monkeypatch, caplog, graph_bs, max_bs):
+    from freetoken.engine.engine import Engine, logger
+    from freetoken.engine.graph import _determine_cuda_graph_bs
+    from freetoken.moe.cpu_executor import CpuMoeExecutor
+
+    torch.manual_seed(0)
+    graph_enabled = bool(_determine_cuda_graph_bs(graph_bs, max_bs, free_memory=1 << 30))
+    monkeypatch.setattr(logger, "propagate", True)
+    original = CpuMoeExecutor.decode_sync
+
+    def corrupt_kernel_capture(self, *args, **kwargs):
+        out = original(self, *args, **kwargs)
+        if self._kernel_io and torch.cuda.is_current_stream_capturing():
+            return out + 1
+        return out
+
+    monkeypatch.setattr(CpuMoeExecutor, "decode_sync", corrupt_kernel_capture)
+    cache = _make_cache(3, 8, 256, 128, scale=0.03)
+    cache.set_cpu_executor = lambda executor: setattr(cache, "cpu_executor", executor)
+    config = SimpleNamespace(
+        max_running_req=1, decode_rows_per_req=1, moe_cpu_threads=4,
+        cuda_graph_bs=graph_bs, cuda_graph_max_bs=max_bs, tp_info=SimpleNamespace(size=1),
+    )
+    layers = [SimpleNamespace(
+        layer_id=i, top_k=4, activation="silu", apply_router_weight_on_input=False,
+        alpha=1.702, limit=None, quant_method=SimpleNamespace(cpu_spec=lambda: ("bf16", {})),
+    ) for i in range(3)]
+    engine = SimpleNamespace(device=torch.device("cuda"), dtype=torch.bfloat16, tp_cpu_group=None)
+    Engine._init_cpu_moe_executor(engine, config, cache, layers)
+    ex = engine.cpu_moe_executor
+    assert "not graph-safe" not in caplog.text
+    assert bool(_determine_cuda_graph_bs(config.cuda_graph_bs, config.cuda_graph_max_bs, free_memory=1 << 30)) == graph_enabled
+    assert ex._flag_sync and ex._kernel_io == (not graph_enabled)
+    assert ex.graph_capture_safe(torch.bfloat16) == graph_enabled
+    hidden = torch.randn(1, 256, device="cuda", dtype=torch.bfloat16)
+    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")
+    if not graph_enabled:
+        actual = ex.decode(1, hidden, weights, ids)
+        expected = _reference(cache, 1, hidden, ids.cpu(), weights.cpu())
+        relative = (actual.float().cpu() - expected).abs().max() / (expected.abs().max() + 1e-6)
+        assert relative < 2e-2, relative
+        return
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        actual = ex.decode(1, hidden, weights, ids)
+    for step in range(3):
+        hidden.add_(0.125)
+        graph.replay()
+        torch.cuda.synchronize()
+        expected = _reference(cache, 1, hidden, ids.cpu(), weights.cpu())
+        relative = (actual.float().cpu() - expected).abs().max() / (expected.abs().max() + 1e-6)
+        assert relative < 2e-2, relative
+
+
+@pytest.mark.skipif(torch.version.hip is None, reason="requires ROCm")
+@pytest.mark.parametrize("kernel_only", [False, True])
+@pytest.mark.parametrize("graph_enabled", [False, True])
+def test_cpu_moe_hip_probe_rejects_common_eager_graph_error(monkeypatch, kernel_only, graph_enabled):
+    from freetoken.moe.cpu_executor import CpuMoeExecutor
+
+    original = CpuMoeExecutor.decode_sync
+
+    def corrupt(self, pending, add_to=None):
+        out = original(self, pending, add_to=add_to)
+        return out + 1 if not kernel_only or pending[2] else out
+
+    monkeypatch.setattr(CpuMoeExecutor, "decode_sync", corrupt)
+    cache = _make_cache(3, 8, 256, 128, scale=0.03)
+
+    def make():
+        return CpuMoeExecutor(cache, top_k=4, activation="silu", apply_router_weight_on_input=False,
+                              num_threads=4, max_tokens=2, device=torch.device("cuda"), graph_enabled=graph_enabled)
+
+    if kernel_only:
+        ex = make()
+        assert not ex._kernel_io
+        assert ex.graph_capture_safe(torch.bfloat16) == graph_enabled
+    else:
+        with pytest.raises(RuntimeError, match="eager IO probe failed"):
+            make()
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+@pytest.mark.parametrize("accumulate", [False, True])
+def test_cpu_moe_graph_dtype_and_hybrid_accumulation(dtype, accumulate):
+    from freetoken.moe.cpu_executor import CpuMoeExecutor
+
+    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=2, device=torch.device("cuda"), activation_dtype=dtype)
+    assert ex.graph_capture_safe(dtype)
+    hidden = torch.randn(2, 256, device="cuda", dtype=dtype)
+    ids = torch.tensor([[0, 2, -1, 5], [-1, -1, -1, -1]], device="cuda", dtype=torch.int32)
+    weights = torch.rand(2, 4, device="cuda")
+    add = torch.randn_like(hidden)
+
+    def forward():
+        pending = ex.decode_submit(1, hidden, weights, ids)
+        assert pending[2] == (ex._kernel_io and dtype == torch.bfloat16)
+        return ex.decode_sync(pending, add_to=add.clone() if accumulate else None)
+
+    forward()
+    torch.cuda.synchronize()
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        actual = forward()
+    for step in range(6):
+        hidden.add_(0.125)
+        ids.fill_(-1) if step % 2 else ids[0].copy_(torch.tensor([1, 3, -1, 7], device="cuda"))
+        graph.replay()
+        torch.cuda.synchronize()
+        ex._io["x"][:2].copy_(hidden.cpu())
+        ex._io["ids"][:2].copy_(ids.cpu())
+        ex._io["w"][:2].copy_(weights.cpu())
+        ex._ext.run(1, 2)
+        expected = ex._io["y"][:2].to(dtype).clone()
+        if accumulate:
+            expected += add.cpu()
+        assert torch.equal(actual.cpu(), expected)
+
+
+@pytest.mark.parametrize("kernel_io", [False, True])
+def test_cpu_moe_hip_event_record_failure_is_loud(hip_probe_fault_library, kernel_io):
     import os
     from pathlib import Path
     import subprocess
     import sys
 
     script = r"""
-import ctypes, runpy, sys
+import ctypes, os, runpy, sys
 import pytest
 import torch
+from freetoken.kernel import _cpu_moe
 from freetoken.moe.cpu_executor import CpuMoeExecutor
 helpers = runpy.run_path(sys.argv[1])
 shim = ctypes.CDLL(sys.argv[2])
+kernel_io = sys.argv[3] == 'True'
+if not kernel_io:
+    _cpu_moe.kernel_io_probe = lambda stream: 'forced memcpy fallback'
 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,
+def make_executor():
+    return CpuMoeExecutor(cache, top_k=4, activation='silu', apply_router_weight_on_input=False,
         num_threads=4, max_tokens=1, device=torch.device('cuda'))
+if kernel_io:
+    ex = make_executor()
+    assert ex._kernel_io and ex.graph_capture_safe(torch.bfloat16)
+    hidden = torch.randn(1, 256, device='cuda', dtype=torch.bfloat16)
+    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')
+    ex.decode(1, hidden, weights, ids)
+    torch.cuda.synchronize()
+    assert shim.probe_test_injected() == 0
+    with pytest.raises(RuntimeError, match='publish input copies'):
+        ex.decode(1, hidden.half(), weights, ids)
+else:
+    with pytest.raises(RuntimeError, match='publish input copies'):
+        make_executor()
 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)],
+        [sys.executable, "-c", script, str(Path(__file__).resolve()), str(hip_probe_fault_library), str(kernel_io)],
         env=env, capture_output=True, text=True, timeout=30,
     )
     assert result.returncode == 0, result.stdout + result.stderr
diff --git a/tests/moe/test_cpu_moe_fp8_block.py b/tests/moe/test_cpu_moe_fp8_block.py
index 8f25b5ee8..17cc5fb57 100644
--- a/tests/moe/test_cpu_moe_fp8_block.py
+++ b/tests/moe/test_cpu_moe_fp8_block.py
@@ -316,16 +316,20 @@ def forward():
 
 
 @pytest.mark.skipif(torch.version.hip is None, reason="requires ROCm")
-def test_cpu_fp8_block_event_graph_replay_and_recapture():
+@pytest.mark.parametrize("kernel_io", [True, False])
+def test_cpu_fp8_block_graph_replay_and_recapture(monkeypatch, kernel_io):
+    from freetoken.kernel import _cpu_moe
     from freetoken.moe.cpu_executor import CpuMoeExecutor
 
+    if not kernel_io:
+        monkeypatch.setattr(_cpu_moe, "kernel_io_probe", lambda stream: "forced memcpy fallback")
     L, E, H, I, top_k = 3, 8, 256, 128, 4
     cache = _make_fp8_block_cache(L, E, H, I, seed=71)
     ex = CpuMoeExecutor(
         cache, top_k=top_k, activation="silu", apply_router_weight_on_input=False,
         num_threads=4, max_tokens=3, device=torch.device("cuda"),
     )
-    assert ex._flag_sync and not ex._kernel_io
+    assert ex._flag_sync and ex._kernel_io == kernel_io
     assert ex.graph_capture_safe(torch.bfloat16)
     stream = torch.cuda.Stream()
     stream.wait_stream(torch.cuda.current_stream())
@@ -348,7 +352,7 @@ def forward():
                 return [ex.decode(layer, hidden[layer], weights[layer], ids[layer]) for layer in range(L)]
 
             inputs(0)
-            # New tasks may allocate their publication event inside capture.
+            # The memcpy fallback may allocate its publication event inside capture.
             graph = torch.cuda.CUDAGraph()
             with torch.cuda.graph(graph, stream=stream):
                 outputs = forward()

From 5e1f3b3fd998f1cb78bfe0609c651e0f652f1224 Mon Sep 17 00:00:00 2001
From: LZ-QWQ <lz1830837889@163.com>
Date: Fri, 30 Oct 2026 18:17:43 +0800
Subject: [PATCH 4/4] fix(moe): enable locked hybrid staging on rocm

---
 docs/install_amd.md                                   | 5 +++++
 python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp | 6 +++---
 python/freetoken/moe/locked_fetch.py                  | 4 ++--
 tests/moe/test_offload.py                             | 8 ++++++--
 4 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/docs/install_amd.md b/docs/install_amd.md
index fed29d0ec..7a228c4ea 100644
--- a/docs/install_amd.md
+++ b/docs/install_amd.md
@@ -75,5 +75,10 @@ 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.
diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp
index c834f4044..ad69f5127 100644
--- a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp
+++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp
@@ -163,8 +163,8 @@ struct GpuSync::Impl {
   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) {
@@ -173,7 +173,7 @@ 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() {
diff --git a/python/freetoken/moe/locked_fetch.py b/python/freetoken/moe/locked_fetch.py
index b2e4331e0..f91770d82 100644
--- a/python/freetoken/moe/locked_fetch.py
+++ b/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:
diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index 3baddeb52..92ea01ee3 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -1129,14 +1129,18 @@ def test_locked_hybrid_layer_fetches_through_staging_in_a_graph(monkeypatch, mem
     # memops: the GPU waits on the doorbell with a stream memop, else the gather kernel spins on it
     from types import SimpleNamespace
 
-    from freetoken.moe.locked_fetch import LockedLayerFetch
+    from freetoken.moe.locked_fetch import attach_locked_fetch, settle_banks
 
     monkeypatch.setenv("FREETOKEN_LOCKED_FETCH_THREADS", "2")
     torch.manual_seed(0)
     cache, sources = _make_split_cache(num_layers=2, locked=(1,), device="cuda")
     cache.decode_target = "hybrid"
     cache.locked_hybrid = frozenset({1})
-    cache.locked_fetch = LockedLayerFetch(cache, [1], plan={1: (0, 1.0)}, max_fetch=2, executor=SimpleNamespace(core_ids=[], _coord_core=-1, _flag_sync=memops))
+    config = SimpleNamespace(moe_strategy="hybrid", moe_locked_fetch=1.0, moe_locked_max_fetch=2)
+    _, _, staging = settle_banks(config, SimpleNamespace(sources=sources), probe=False, has_locked=True)
+    assert staging is not None
+    cache.cpu_executor = SimpleNamespace(core_ids=[], _coord_core=-1, _flag_sync=memops)
+    attach_locked_fetch(cache, config, staging, layer_rows=None)
 
     top_k = 3
     ids = torch.zeros(top_k, dtype=torch.int32, device="cuda")